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

from pathlib import Path

import tripleblind as tb


DATA_OWNER = tb.config.example_user1
CONSUMER = tb.config.example_user2
data_dir = Path("example_data")

tb.initialize(api_token=DATA_OWNER["token"], example=True)

try:
    print("Creating dataset")

    asset0 = tb.Asset.position(
        file_handle=data_dir / "train.zip",
        name="Energy forecast training data",
        desc="Energy forecast training data",
        is_discoverable=False,  # keep these datasets private between the teams
    )

    if asset0 is not None:
        print("Created new dataset")

    # The team which owns the data must create an agreement with the
    # data consuming team in order for them to access it in training and
    # inference.
    asset0.add_agreement(
        with_team=CONSUMER["team_id"],
        operation=tb.Operation.FEDERATED_LEARNING,
    )
    print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset already exists.")
print("Training data is in position.")


asset1 = None
try:
    print("Creating dataset")

    asset1 = tb.Asset.position(
        file_handle=data_dir / "test.zip",
        name="Energy forecast testing data",
        desc="Energy forecast testing data",
        is_discoverable=False,
    )

    if asset1 is not None:
        print("Created new dataset")

    asset1.add_agreement(
        with_team=CONSUMER["team_id"],
        operation=tb.Operation.FEDERATED_LEARNING,
    )
    print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset already exists.")

print("Testing data is in position.")
