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

import tripleblind as tb


##########################################################################
# GET AUTHENTICATION TOKENS AND ESTABLISH CONNECTION TO THE ROUTER
#
# Establish the connection details to reach the TripleBlind instance.
# Unless explicitly specified, all operations will occur via this default
# session as the user 'organization_one'
tb.initialize(api_token=tb.config.example_user1["token"])

name1 = "EXAMPLE - NASA Turbofan Failure Data (Part 1)"  # org 3
name2 = "EXAMPLE - NASA Turbofan Failure Data (Part 2)"  # org 2
name3 = "EXAMPLE - NASA Turbofan Failure Data (Part 3)"  # org 1
name4 = "EXAMPLE - NASA Turbofan Failure Data (Part 4)"  # org 3
desc = """The NASA C-MAPSS (Commercial Modular Aero-Propulsion System
Simulation) Aircraft Engine dataset, containing engine sensor data over time
until the turbofan engine reaches failure. This dataset has been split into
four parts to illustrate working with information distributed across different
organizations.

Original dataset can be found on the [NASA Intelligent Systems Division website]
(https://ti.arc.nasa.gov/tech/dash/groups/pcoe/prognostic-data-repository/#turbofan)."""

name_nn = "EXAMPLE - NASA Turbofan Failure Data (with Target)"  # org 2
desc_nn = """The NASA C-MAPSS (Commercial Modular Aero-Propulsion System
Simulation) Aircraft Engine dataset, containing engine sensor data over time
until the turbofan engine reaches failure. This dataset includes a Target column
for use in neural network training.

Original dataset can be found on the [NASA Intelligent Systems Division website]
(https://ti.arc.nasa.gov/tech/dash/groups/pcoe/prognostic-data-repository/#turbofan)."""

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

    print(f"Creating '{name1}' on {tb.config.example_user3['login']}'s Access Point...")
    dataset_train1 = tb.Asset.position(
        file_handle="FD001_train.zip",
        name=name1,
        desc=desc,
        is_discoverable=True,
        session=session3,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{name1}' already exists, skipping'")


try:
    # Start a connection as 'organization_two' user
    session2 = tb.Session(api_token=tb.config.example_user2["token"], from_default=True)

    print(f"Creating '{name2}' on {tb.config.example_user2['login']}'s Access Point...")
    dataset_train2 = tb.Asset.position(
        file_handle="FD002_train.zip",
        name=name2,
        desc=desc,
        is_discoverable=True,
        session=session2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{name2}' already exists, skipping'")
    dataset_train2 = tb.Asset.find(name2)
if tb.config.create_agreements_on_example_input_assets:
    # For the examples we will attach an Agreement to the datasets.  This
    # agreement makes the dataset available to the other organization,
    # meaning the training step will not require the dataset owner to
    # explicitly grant permission for use.
    dataset_train2.add_agreement(
        with_team=tb.config.example_user3["team_id"],
        operation=tb.Operation.BLIND_LEARNING,
        session=session2,
    )
    print("Created Agreement")


try:
    print(
        f"Creating '{name_nn}' on {tb.config.example_user2['login']}'s Access Point..."
    )
    dataset_nn = tb.Asset.position(
        file_handle="example_data/FD001_nn_train.csv",
        name=name_nn,
        desc=desc_nn,
        is_discoverable=True,
        session=session2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{name_nn}' already exists, skipping'")
    dataset_nn = tb.Asset.find(name_nn)
if tb.config.create_agreements_on_example_input_assets:
    dataset_nn.add_agreement(
        with_team=tb.config.example_user3["team_id"],
        operation=tb.Operation.BLIND_LEARNING,
        session=session2,
    )
    print("Created Agreement")


try:
    # Start a connection as 'organization_one' user
    session1 = tb.Session(api_token=tb.config.example_user1["token"], from_default=True)

    print(f"Creating '{name3}' on {tb.config.example_user1['login']}'s Access Point...")
    dataset_train3 = tb.Asset.position(
        file_handle="FD003_train.zip",
        name=name3,
        desc=desc,
        is_discoverable=True,
        session=session1,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{name3}' already exists, skipping'")
    dataset_train3 = tb.Asset.find(name3)
if tb.config.create_agreements_on_example_input_assets:
    dataset_train3.add_agreement(
        with_team=tb.config.example_user3["team_id"],
        operation=tb.Operation.BLIND_LEARNING,
        session=session1,
    )
    print("Created Agreement")


try:
    print(f"Creating '{name4}' on {tb.config.example_user3['login']}'s Access Point...")
    dataset_train4 = tb.Asset.position(
        file_handle="FD004_train.zip",
        name=name4,
        desc=desc,
        is_discoverable=True,
        session=session3,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{name4}' already exists, skipping'")


print("Data is in position.")
