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

import tripleblind as tb


# Establish a connection to the Router as `example_user1`
tb.util.set_script_dir_current()
tb.initialize(api_token=tb.config.example_user1["token"])

name = "EXAMPLE - Movie rating database - part one"
try:
    # Register the training database with the Router and place it the AP
    print("Uploading example dataset to Org 1's access point...")
    asset = tb.asset.CSVDataset.create(
        datafile="movie_rating_train_part_one.csv",
        name=name,
        desc="""Synthetic dataset for used in the NLP/Sequence Classification
example. The dataset holds free-text movie reviews along with a numeric (1 for positive sentiment or 0 for negative) ratings.

Part one of two distributed datasets.""",
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    # If the dataset already exists, retrieve it
    print(f"   asset '{name}' already exists, skipping.")
    asset = tb.Asset.find(name, owned_by=tb.config.example_user1["team_id"])

# Attach an Agreement to the dataset to make it available to
# `example_user3` explicitly to run the SC_TRAIN protocol only.
asset.add_agreement(
    with_team=tb.config.example_user3["team_id"],
    operation=tb.Operation.SC_TRAIN,
)

tb.initialize(api_token=tb.config.example_user2["token"])

name = "EXAMPLE - Movie rating database - part two"
try:
    # Register the training database with the Router and place it the AP
    print("Uploading example dataset to Org 2's access point...")
    asset = tb.asset.CSVDataset.create(
        datafile="movie_rating_train_part_two.csv",
        name=name,
        desc="""Synthetic dataset for used in the NLP/Sequence Classification
example. The dataset holds free-text movie reviews along with a numeric (1 for positive sentiment or 0 for negative) ratings.

Part two of two distributed datasets.""",
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    # If the dataset already exists, retrieve it
    print(f"   asset '{name}' already exists, skipping.")
    asset = tb.Asset.find(name, owned_by=tb.config.example_user2["team_id"])

# Attach an Agreement to the dataset to make it available to
# `example_user3` explicitly to run the SC_TRAIN protocol only.
asset.add_agreement(
    with_team=tb.config.example_user3["team_id"],
    operation=tb.Operation.SC_TRAIN,
)
