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

from pathlib import Path

import tripleblind as tb


user1 = tb.config.example_user1
user2 = tb.config.example_user2

name1 = "EXAMPLE - Experian Credit Score Data"
desc1 = """
This is a fictional dataset that contains personal financial data that can be
used to determine an individual's credit score. The dataset contains a set of
features and a target value. The target value is the change in an individual's
credit score over a fixed period of time.
See associated dataset EXAMPLE - TransUnion Credit Score Data.
"""

name2 = "EXAMPLE - TransUnion Credit Score Data"
desc2 = """
This is a fictional dataset that contains personal financial data that can be
used to determine an individual's credit score. The dataset contains a set of
features and a target value. The target value is the change in an individual's
credit score over a fixed period of time.
See associated dataset EXAMPLE - Experian Credit Score Data.
"""

##############################################################################
# FOR BREVITY, WE WILL UPLOAD DATA FOR ALL THREE COMPANIES IN ONE JOB.
# IN A REAL SITUATION THESE WOULD BE INDEPENDENT JOBS RUN BY THE SEPARATE
# ORGANIZATIONS

tb.util.set_script_dir_current()
data_dir = Path("example_data")
data_dir.mkdir(exist_ok=True)

##############################################################################
#  Company one
#############################################################################

# Upload the first training dataset to Company one's Access Point
try:
    # Start a connection as 'organization_one' user
    session1 = tb.Session(api_token=tb.config.example_user1["token"], from_default=True)

    # Make data available on the Router index. Position it on the Access Point
    # and save meta-data about the dataset in the Router index.
    #
    print(f"Creating dataset '{name1}' on {user1['login']}'s Access Point...")
    dataset_train0 = tb.asset.CSVDataset.position(
        file_handle=data_dir / "split_0_reg_random_train.csv",
        name=name1,
        desc=desc1,
        is_discoverable=True,
        session=session1,
        auto_rename_columns=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"   asset '{name1}' already exists, skipping upload.")
    dataset_train0 = tb.Asset.find(name1, owned_by=tb.config.example_user1["team_id"])

# Attach an Agreement to the dataset to allow anyone to train against
# this dataset without further interaction.
dataset_train0.add_agreement(
    with_team="ANY", operation=tb.Operation.XGBOOST_TRAIN, session=session1
)

#############################################################################
#  Company two
#############################################################################

# Upload organization-two dataset to Company two's Access Point
try:
    # Start a connection as 'organization_two' user
    session2 = tb.Session(api_token=tb.config.example_user2["token"], from_default=True)

    # Make data available on the Router index.  Position it on the Access Point
    # and save meta-data about the dataset in the Router index.
    #
    print(f"Creating dataset '{name2}' on {user2['login']}'s Access Point...")
    dataset_train1 = tb.asset.CSVDataset.position(
        file_handle=data_dir / "split_1_reg_random_train.csv",
        name=name2,
        desc=desc2,
        is_discoverable=True,
        session=session2,
        auto_rename_columns=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"   asset '{name2}' already exists, skipping upload.")
    dataset_train1 = tb.Asset.find(name2, owned_by=tb.config.example_user2["team_id"])

# Attach an Agreement to the dataset to allow anyone to train against
# this dataset without further interaction.
dataset_train1.add_agreement(
    with_team="ANY", operation=tb.Operation.XGBOOST_TRAIN, session=session2
)

print("Data is in position.")
