#!/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_user2`
tb.util.set_script_dir_current()
tb.initialize(api_token=tb.config.example_user2["token"])

name = "EXAMPLE - Natural Language Processing training data"
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="synthetic_data.csv",
        name=name,
        desc="Synthetic dataset for medical Named Entity Recognition (NER).",
        is_discoverable=True,
        allow_overwrite=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 NLP_TRAIN protocol only.
asset.add_agreement(
    with_team=tb.config.example_user3["team_id"],
    operation=tb.Operation.NLP_TRAIN,
)
