#!/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 - "

people = dict(
    alice=tb.config.example_user1["token"],
    bob=tb.config.example_user2["token"],
)

for person, token in people.items():
    tb.initialize(api_token=token)

    # Position training data on access point.
    training_asset = tb.Asset.position(
        file_handle=f"{data_dir}/{person}_training.csv",
        name=f"{prefix}{person}'s portion of a PSI Vertically Partitioned dataset for training.",
        desc="One party's portion of a PSI Vertically Partitioned dataset.",
        allow_overwrite=True,
        is_discoverable=True,
    )

    # Add agreement to allow K Means Clustering training with training data.
    training_asset.add_agreement(
        with_team="ANY",
        operation=tb.Operation.PSI_VERTICAL_KMEANS_TRAIN,
    )

    # Position inference data on access point.
    inference_asset = tb.Asset.position(
        file_handle=f"{data_dir}/{person}_inference.csv",
        name=f"{prefix}{person}'s portion of a PSI Vertically Partitioned dataset for inference.",
        desc="One party's portion of a PSI Vertically Partitioned dataset.",
        allow_overwrite=True,
        is_discoverable=True,
    )
