#!/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 test 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

name1 = "TEST - CIFAR-10, part 1"
desc1 = """
Small version (10 image) of the CIFAR-10 dataset for rapid testing.
"""

name2 = "TEST - CIFAR-10, part 2"
desc2 = """
Small version (10 image) of the CIFAR-10 dataset for rapid testing.
"""


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

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 / "TEST_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 / "TEST_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.")
