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

import tripleblind as tb


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

user1 = tb.config.example_user1
user3 = tb.config.example_user3

run_id = tb.util.read_run_id()

if os.environ.get("TB_TEST_SMALL"):
    model_file = "iris-one.xml"
    name = f"TEST - PMML Random Forest model - {run_id}"
else:
    model_file = "iris-randomForest-R.xml"
    name = "EXAMPLE - PMML Random Forest model"

description = """Simple Predictive Model Markup Language model for predicting
flower species from flower characteristics."""


print(f"Positioning '{model_file}' as {name} on {user1['login']}'s Access Point...")
try:
    model = tb.asset.PMMLTree.create(
        model_file,
        name=name,
        desc=description,
        allow_overwrite=True,
    )
except tb.TripleblindAssetAlreadyExists:
    print(f"   Asset '{name}' already exists.'")
    try:
        model = tb.Asset.find(name, owned_by=user1["team_id"])
    except Exception as e:
        print("Could not find model using tb.Asset.find")

# Attach an Agreement to the algorithm.  This agreement makes the algorithm
# available to the other team, which means the inference step will not
# require the algorithm owner to explicitly grant permission for use.  Usage is
# still logged for auditing.
model.add_agreement(with_team=user3["team_id"], operation=tb.Operation.EXECUTE)
print(f"Created Agreement with '{user3['login']}'")
