#!/usr/bin/env python
# Copyright (c) TripleBlind Holdings, Inc. Confidential and Proprietary. All rights reserved.

from pathlib import Path

import tripleblind as tb


##############################################################################
#
# Script to create sample CIFAR datasets.  Two datasets are created and
# placed on the access points for example_user1, and user2.  Permissions are
# set for training jobs to operate hands-free.
#
##############################################################################

user1 = tb.config.example_user1
user2 = tb.config.example_user2
user3 = tb.config.example_user3

name1 = "EXAMPLE - CIFAR-10, part 1"
desc1 = """
Half of the well-known CIFAR-10 dataset of color images.  This dataset of small
(32x32) images can be used for image classification problems.  Traditionally
it is used as a single dataset for training, but here it has been chopped in
half to illustrate distributed training.

Images are in one of ten different classes:

|   | Class        | Sample
|---|--------------|--------
|   | _airplane_   | ![sample plane](https://www.cs.toronto.edu/~kriz/cifar-10-sample/airplane4.png)
|   | _automobile_ | ![sample car](https://www.cs.toronto.edu/~kriz/cifar-10-sample/automobile10.png)
|   | _bird_       | ![sample bird](https://www.cs.toronto.edu/~kriz/cifar-10-sample/bird6.png)
|   | _cat_        | ![sample cat](https://www.cs.toronto.edu/~kriz/cifar-10-sample/cat9.png)
|   | _horse_      | ![sample horse](https://www.cs.toronto.edu/~kriz/cifar-10-sample/horse6.png)
|   | _deer_       | ![sample deer](https://www.cs.toronto.edu/~kriz/cifar-10-sample/deer4.png)
|   | _dog_        | ![sample dog](https://www.cs.toronto.edu/~kriz/cifar-10-sample/dog4.png)
|   | _frog_       | ![sample frog](https://www.cs.toronto.edu/~kriz/cifar-10-sample/frog6.png)
|   | _ship_       |  ![sample ship](https://www.cs.toronto.edu/~kriz/cifar-10-sample/ship1.png)
|   | _truck_      | ![sample truck](https://www.cs.toronto.edu/~kriz/cifar-10-sample/truck5.png)

Classes are mutually exclusive -- a "truck" is never an "automobile".

More information on the Canadian Institute for Advanced Research (CIFAR) and
this dataset can be found at [Alex Krizhevsky's CIFAR page](https://www.cs.toronto.edu/~kriz/cifar.html).
"""

name2 = "EXAMPLE - CIFAR-10, part 2"
desc2 = """
Half of the well-known CIFAR-10 dataset of color images.  This dataset of small
(32x32) images can be used for image classification problems.  Traditionally
it is used as a single dataset for training, but here it has been chopped in
half to illustrate distributed training.

Images are in one of ten different classes:

|   | Class        | Sample
|---|--------------|--------
|   | _airplane_   | ![sample plane](https://www.cs.toronto.edu/~kriz/cifar-10-sample/airplane5.png)
|   | _automobile_ | ![sample car](https://www.cs.toronto.edu/~kriz/cifar-10-sample/automobile4.png)
|   | _bird_       | ![sample bird](https://www.cs.toronto.edu/~kriz/cifar-10-sample/bird2.png)
|   | _cat_        | ![sample cat](https://www.cs.toronto.edu/~kriz/cifar-10-sample/cat8.png)
|   | _horse_      | ![sample horse](https://www.cs.toronto.edu/~kriz/cifar-10-sample/horse1.png)
|   | _deer_       | ![sample deer](https://www.cs.toronto.edu/~kriz/cifar-10-sample/deer2.png)
|   | _dog_        | ![sample dog](https://www.cs.toronto.edu/~kriz/cifar-10-sample/dog3.png)
|   | _frog_       | ![sample frog](https://www.cs.toronto.edu/~kriz/cifar-10-sample/frog1.png)
|   | _ship_       |  ![sample ship](https://www.cs.toronto.edu/~kriz/cifar-10-sample/ship7.png)
|   | _truck_      | ![sample truck](https://www.cs.toronto.edu/~kriz/cifar-10-sample/truck8.png)

Classes are mutually exclusive -- a "truck" is never an "automobile".

More information on the Canadian Institute for Advanced Research (CIFAR) and
this dataset can be found at [Alex Krizhevsky's CIFAR page](https://www.cs.toronto.edu/~kriz/cifar.html).
"""


##########################################################################

tb.initialize(api_token=user1["token"])
tb.util.set_script_dir_current()
data_dir = Path("example_data")


#############################################################################
# Register the first training database with the Router and place on Access Point
#############################################################################

print(
    f"Creating dataset '{name1}' on {tb.config.example_user1['login']}'s Access Point..."
)
try:
    dataset1 = tb.Asset.position(
        file_handle=tb.Package.load(data_dir / "EXAMPLE_train_cifar_0.zip"),
        name=name1,
        desc=desc1,
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{name1}' already exists.'")
    dataset1 = tb.Asset.find(name1, owned_by=tb.config.example_user1["team_id"])

# Attach an Agreement to the datasets.  This agreement makes the dataset
# available to the training organization of user3.
dataset1.add_agreement(with_team="ANY", operation=tb.Operation.BLIND_LEARNING)
print("Created Agreement for any to train against this dataset.\n")

#############################################################################
# Register the second training database with the Router and place on Access Point
#############################################################################

session2 = tb.Session(api_token=user2["token"], from_default=True)

print(f"Creating dataset '{name2}' on {user2['login']}'s Access Point...")
try:
    dataset2 = tb.Asset.position(
        file_handle=tb.Package.load(data_dir / "EXAMPLE_train_cifar_1.zip"),
        name=name2,
        desc=desc2,
        is_discoverable=True,
        session=session2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{name2}' already exists.'")
    dataset2 = tb.Asset.find(name2, owned_by=tb.config.example_user2["team_id"])

# Attach an Agreement to the datasets.  This agreement makes the dataset
# available to the training organization of user3.
dataset2.add_agreement(
    with_team="ANY", operation=tb.Operation.BLIND_LEARNING, session=session2
)
print("Created Agreement for any to train against this dataset.\n")


print("Data is in position.")
