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

from pathlib import Path

import tripleblind as tb


tb.util.set_script_dir_current()
data_dir = Path("example_data")

prefix = "EXAMPLE - "
train0 = "Decision Tree classification, training data (part 1)"
train1 = "Decision Tree classification, training data (part 2)"
train2 = "Decision Tree classification, training data (part 3)"
test0 = "Decision Tree classification, test data (part 1)"
test1 = "Decision Tree classification, test data (part 2)"
test2 = "Decision Tree classification, test data (part 3)"
test_all = "Decision Tree classification, test truth"

reg_train0 = "Decision Tree regression, training data (part 1)"
reg_train1 = "Decision Tree regression, training data (part 2)"
reg_test0 = "Decision Tree regression, test data (part 1)"
reg_test1 = "Decision Tree regression, test data (part 2)"
reg_test_all = "Decision Tree regression, test truth"


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

##########################################################################
# Place data on the organization-one Access Point

try:
    print(f"Uploading '{prefix}{train0}' to organization-one's Access Point...")
    train_0 = tb.Asset.position(
        data_dir / "train_clf0.csv",
        name=f"{prefix}{train0}",
        desc="""The first part of a distributed training dataset used in the
        Decision_Tree example. This portion contains 33% of the data.""",
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{train0}' already exists'")


try:
    print(f"Uploading '{prefix}{test1}' to organization-one's Access Point...")
    test_0 = tb.Asset.position(
        data_dir / "test_clf0.csv",
        # was name=f"PSI DT Classification vertical test 0-40-{run_id}",
        name=f"{prefix}{test0}",
        desc="""The first part of a distributed testing dataset used in the
        Decision_Tree example. This portion contains 33% of the data.""",
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{test0}' already exists'")


try:
    print(f"Uploading '{prefix}{test_all}' to organization-one's Access Point...")
    test_0 = tb.Asset.position(
        data_dir / "test_clf.csv",
        name=f"{prefix}{test_all}",
        desc="""Stash the complete test dataset to be used for accuracy
        comparison in the classification example.""",
        is_discoverable=False,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{test_all}' already exists'")

##########################################################################
# Place data on the organization-two Access Point

try:
    print(f"Uploading '{prefix}{train1}' to organization-two's Access Point...")
    train_1 = tb.Asset.position(
        data_dir / "train_clf1.csv",
        # was name=f"PSI DT Classification vertical train 41-100-{run_id}",
        name=f"{prefix}{train1}",
        desc="""The second part of a distributed training dataset used in the
        Decision_Tree example. This portion contains 50% of the data.""",
        is_discoverable=True,
        session=session_org_2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{train1}' already exists'")
    train_1 = tb.Asset.find(
        f"{prefix}{train1}", owned_by=tb.config.example_user2["team_id"]
    )

train_1.add_agreement(
    with_team=tb.config.example_user1["team_id"],
    operation=tb.Operation.PSI_VERTICAL_DECISION_TREE_TRAIN,
    session=session_org_2,
)
print("Added agreement for use by org1")


try:
    print(f"Uploading '{prefix}{test1}' to organization-two's Access Point...")
    test_1 = tb.Asset.position(
        data_dir / "test_clf1.csv",
        # was name=f"PSI DT Classification vertical test 41-100-{run_id}",
        name=f"{prefix}{test1}",
        desc="""The second part of a distributed testing dataset used in the
        Decision_Tree example. This portion contains 50% of the data.""",
        is_discoverable=True,
        session=session_org_2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{test0}' already exists'")


##########################################################################
# Place data on the organization-three Access Point

try:
    print(f"Uploading '{prefix}{train2}' to organization-three's Access Point...")
    train_2 = tb.Asset.position(
        data_dir / "train_clf2.csv",
        name=f"{prefix}{train2}",
        # name=f"PSI DT Classification vertical train 101-120-{run_id}",
        desc="""The third part of a distributed training dataset used in the
        Decision_Tree example. This portion contains 17% of the data.""",
        is_discoverable=True,
        session=session_org_3,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{train2}' already exists'")
    train_2 = tb.Asset.find(
        f"{prefix}{train2}", owned_by=tb.config.example_user3["team_id"]
    )

train_2.add_agreement(
    with_team=tb.config.example_user1["team_id"],
    operation=tb.Operation.PSI_VERTICAL_DECISION_TREE_TRAIN,
    session=session_org_3,
)
print("Added agreement for use by org1")


try:
    print(f"Uploading '{prefix}{test2}' to organization-three's Access Point...")
    test_2 = tb.Asset.position(
        data_dir / "test_clf2.csv",
        # was name=f"PSI DT Classification vertical test 101-120-{run_id}",
        name=f"{prefix}{test2}",
        desc="""The third part of a distributed testing dataset used in the
        Decision_Tree example. This portion contains 17% of the data.""",
        is_discoverable=True,
        session=session_org_3,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{test2}' already exists'")

print("Classification data is in position.")


##########################################################################
# Regression Data
##########################################################################

try:
    print(f"Uploading '{prefix}{reg_train0}' to organization-one's Access Point...")
    reg_train_0 = tb.Asset.position(
        data_dir / "train_reg0.csv",
        # was name=f"PSI DT Regression vertical train 0-40-{run_id}",
        name=f"{prefix}{reg_train0}",
        desc="""The first half of a distributed training dataset used in the
        Decision_Tree regression example.""",
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{reg_train0}' already exists'")

try:
    print(f"Uploading '{prefix}{reg_test0}' to organization-one's Access Point...")
    reg_test_0 = tb.Asset.position(
        data_dir / "test_reg0.csv",
        # was name=f"PSI DT Regression vertical test 0-40-{run_id}",
        name=f"{prefix}{reg_test0}",
        desc="""The first half of a distributed testing dataset used in the
        Decision_Tree regression example.""",
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{reg_test0}' already exists'")

try:
    print(f"Uploading '{prefix}{reg_test_all}' to organization-one's Access Point...")
    test_0 = tb.Asset.position(
        data_dir / "test_reg.csv",
        name=f"{prefix}{reg_test_all}",
        desc="""Stash the complete test dataset to be used for accuracy
        comparison in the regression example.""",
        is_discoverable=False,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{reg_test_all}' already exists'")


try:
    print(f"Uploading '{prefix}{reg_train1}' to organization-two's Access Point...")
    reg_train_1 = tb.Asset.position(
        data_dir / "train_reg1.csv",
        # was name=f"PSI DT Regression vertical train 41-120-{run_id}",
        name=f"{prefix}{reg_train1}",
        desc="""The second half of a distributed training dataset used in the
        Decision_Tree regression example.""",
        is_discoverable=True,
        session=session_org_2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{reg_train1}' already exists'")
    reg_train_1 = tb.Asset.find(
        f"{prefix}{reg_train1}", owned_by=tb.config.example_user2["team_id"]
    )

reg_train_1.add_agreement(
    with_team=tb.config.example_user1["team_id"],
    operation=tb.Operation.PSI_VERTICAL_DECISION_TREE_TRAIN,
    session=session_org_2,
)
print("Added agreement for use by org1")


try:
    print(f"Uploading '{prefix}{reg_test1}' to organization-two's Access Point...")
    reg_test_1 = tb.Asset.position(
        data_dir / "test_reg1.csv",
        # was name=f"PSI DT Regression vertical test 41-120-{run_id}",
        name=f"{prefix}{reg_test1}",
        desc="""The second half of a distributed testing dataset used in the
        Decision_Tree regression example.""",
        is_discoverable=True,
        session=session_org_2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{prefix}{reg_test1}' already exists'")


print("Regression data  is in position.")
