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

import tripleblind as tb


namet0 = "EXAMPLE - PSI VP XGBoost Training (part 1)"
namet1 = "EXAMPLE - PSI VP XGBoost Training (part 2)"
namei0 = "EXAMPLE - PSI VP XGBoost Inference (part 1)"
namei1 = "EXAMPLE - PSI VP XGBoost Inference (part 2)"
desc = """
Example data for PSI VP XGBoost.
"""


tb.util.set_script_dir_current()
data_dir = Path("example_data")
tb.initialize(api_token=tb.config.example_user1["token"])

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

    print(
        f"Creating '{namet0}' on {tb.config.example_user1['login']}'s Access Point..."
    )
    datasett0 = tb.Asset.position(
        file_handle=data_dir / "datasett0.csv",
        name=namet0,
        desc=desc,
        is_discoverable=True,
        k_grouping=5,
        session=session1,
    )

    print(
        f"Creating '{namet0}' on {tb.config.example_user1['login']}'s Access Point..."
    )
    datasett0 = tb.Asset.position(
        file_handle=data_dir / "datasett0.csv",
        name=namet0,
        desc=desc,
        is_discoverable=True,
        k_grouping=5,
        session=session1,
    )

except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{namet0}' already exists, skipping.")


# Register the second dataset with the Router and place on Access Point
try:
    # Start a connection as 'organization_two' user
    session2 = tb.Session(api_token=tb.config.example_user2["token"], from_default=True)

    print(
        f"Creating '{namet1}' on {tb.config.example_user2['login']}'s Access Point..."
    )
    dataseti1 = tb.Asset.position(
        file_handle=data_dir / "dataseti1.csv",
        name=namet1,
        desc=desc,
        is_discoverable=True,
        session=session2,
    )
    if tb.config.create_agreements_on_example_input_assets:
        # Give user1 access to this file for the PSI_VP_XGBOOST_TRAIN calculation
        dataseti1.add_agreement(
            with_team=tb.config.example_user1["team_id"],
            operation=tb.Operation.PSI_VP_XGBOOST_TRAIN,
            session=session2,
        )
        print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{namet1}' already exists, skipping.")

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

    print(
        f"Creating '{namei0}' on {tb.config.example_user1['login']}'s Access Point..."
    )
    dataseti0 = tb.Asset.position(
        file_handle=data_dir / "dataseti0.csv",
        name=namei0,
        desc=desc,
        is_discoverable=True,
        k_grouping=5,
        session=session1,
    )

    print(
        f"Creating '{namei0}' on {tb.config.example_user1['login']}'s Access Point..."
    )
    dataseti0 = tb.Asset.position(
        file_handle=data_dir / "dataseti0.csv",
        name=namei0,
        desc=desc,
        is_discoverable=True,
        k_grouping=5,
        session=session1,
    )

except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{namei0}' already exists, skipping.")


# Register the second dataset with the Router and place on Access Point
try:
    # Start a connection as 'organization_two' user
    session2 = tb.Session(api_token=tb.config.example_user2["token"], from_default=True)

    print(
        f"Creating '{namei1}' on {tb.config.example_user2['login']}'s Access Point..."
    )
    dataseti1 = tb.Asset.position(
        file_handle=data_dir / "dataseti1.csv",
        name=namei1,
        desc=desc,
        is_discoverable=True,
        session=session2,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{namei1}' already exists, skipping.")

print("Data is in position.")
