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

from pathlib import Path

import tripleblind as tb


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

name1 = "EXAMPLE - Patient Data (imperial)"
name2 = "EXAMPLE - Patient Data (metric)"

try:
    print(f"Creating '{name1}' on {tb.config.example_user1['name']}'s access point...")
    asset_1 = tb.Asset.position(
        file_handle=data_dir / "data_munge_a.csv",
        name=name1,
        desc="Example dataset containing patient information in imperial units.",
        unmask_columns=["Height_IN", "Weight_LBS"],
        is_discoverable=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"'{name1}' already exists, skipping.")


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

try:
    print(f"Creating '{name2}' on {tb.config.example_user2['name']}'s access point...")
    asset_2 = tb.Asset.position(
        file_handle=data_dir / "data_munge_b.csv",
        name=name2,
        desc="Example dataset containing information in metric units on patients over 50 years old.",
        is_discoverable=True,
        session=session_org_2,
    )

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

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 organization.
    asset_2 = tb.Asset.find(name2, owned_by=tb.config.example_user2["team_id"])
    asset_2.add_agreement(
        with_team=tb.config.example_user1["team_id"],
        operation=tb.Operation.REGRESSION,
        session=session_org_2,
    )
    print("Created agreement")

print("Data is in position.")
