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

import tripleblind as tb


run_id = tb.util.read_run_id()

tb.initialize(api_token=tb.config.example_user1["token"])
session_org_2 = tb.Session(
    api_token=tb.config.example_user2["token"], from_default=True
)

try:
    print(
        f"Uploading 'Random Multimodal' to {tb.config.example_user2['login']}'s Access Point..."
    )
    asset = tb.asset.CSVDataset.position(
        tb.config.data_dir / "random_multimodal_train.csv",
        name=f"Random Multimodal-{run_id}",
        desc="Random Multimodal dataset to use with dicom sl example.",
        is_discoverable=True,
        session=session_org_2,
        auto_rename_columns=True,
    )

    if tb.config.create_agreements_on_example_input_assets:
        # For this example we will attach an Agreement to the datasets.  This
        # agreement makes the dataset available to the other team,
        # meaning the training step will not require the dataset owner to
        # explicitly grant permission when it is used in the next step.
        asset.add_agreement(
            with_team=tb.config.example_user1["team_id"],
            operation=tb.Operation.VERTICAL_BLIND_LEARNING,
            session=session_org_2,
        )
        print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print(
        f"Asset 'Random Multimodal 0-{run_id}' already exists, change value in run_id.txt to make it unique."
    )

# Register the second training database with the Router and place on Access Point
try:
    # Start a connection as 'organization_three' user
    session3 = tb.Session(api_token=tb.config.example_user3["token"], from_default=True)

    print(
        f"Uploading 'train_ct_multimodal' to {tb.config.example_user3['login']}'s Access Point..."
    )
    asset = tb.Asset.position(
        file_handle=tb.Package.load(tb.config.data_dir / "train_ct_multimodal.zip"),
        name=f"train_ct_multimodal-{run_id}",
        desc="Dicom images for training multimodal model.",
        is_discoverable=True,
        session=session3,
    )

    if tb.config.create_agreements_on_example_input_assets:
        asset.add_agreement(
            with_team=tb.config.example_user1["team_id"],
            operation=tb.Operation.VERTICAL_BLIND_LEARNING,
            session=session3,
        )
        print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print(
        f"Asset 'train_ct_multimodal-{run_id}' already exists, change value in run_id.txt to make it unique.'"
    )


try:
    print(
        f"Uploading 'random-multimodal-test' to {tb.config.example_user2['login']}'s Access Point..."
    )
    asset = tb.asset.CSVDataset.position(
        tb.config.data_dir / "random_multimodal_test.csv",
        name=f"random-multimodal-test-{run_id}",
        desc="Sample of columns of random set for distributed inference.",
        is_discoverable=True,
        session=session_org_2,
        auto_rename_columns=True,
    )

except tb.TripleblindAssetAlreadyExists:
    print(
        f"Asset 'random-multimodal-test-{run_id}' already exists, change the value in run_id.txt to make it unique."
    )


try:
    print(
        f"Uploading 'test_ct_multimodal' to {tb.config.example_user3['login']}'s Access Point..."
    )
    asset = tb.Asset.position(
        tb.config.data_dir / "test_ct_multimodal.zip",
        name=f"test_ct_multimodal-{run_id}",
        desc="Sample of ct dataset for multimodal inference",
        is_discoverable=True,
        session=session3,
    )

except tb.TripleblindAssetAlreadyExists:
    print(
        f"Asset 'test_ct_multimodal 0-{run_id}' already exists, change the value in run_id.txt to make it unique."
    )


print("Data is in position.")
