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

import tripleblind as tb

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

asset_id = tb.util.load_from("trained_model_asset_id.out")
model = tb.ModelAsset(asset_id)

table1 = tb.TableAsset.find(
    "EXAMPLE - PSI VP XGBoost Inference (part 1)",
    owned_by=tb.config.example_user1["team_id"],
)
table2 = tb.TableAsset.find(
    "EXAMPLE - PSI VP XGBoost Inference (part 2)",
    owned_by=tb.config.example_user2["team_id"],
)

# Create agreements to allow user2 to access both the model and data belonging
# to user1 without requiring manual approval.
agreement_session = tb.Session(
    api_token=tb.config.example_user1["token"], from_default=True
)
model.add_agreement(
    with_team=tb.config.example_user2["team_id"],
    # operation=tb.Operation.PSI_VP_XGBOOST_INFER,
    operation=model,
    session=agreement_session,
)
table1.add_agreement(
    with_team=tb.config.example_user2["team_id"],
    operation=asset_id,
    session=agreement_session,
)


# Use the previously trained XGBoost model to infer against the two
# virtually-joined tables.  This inference is being run by the second
# user, not the owner of the trained model.
result = model.infer(
    data=[table2, table1],
    preprocessor=tb.TabularPreprocessor.builder().all_columns(),
    params={"match_column": ["id", "identifier"]},
)
print(result.dataframe)
result.asset.retrieve(save_as="opposite_inference_output.zip", overwrite=True)
