#!/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_user1["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"],
)

# For this example, add an agreement to avoid a request for approval when
# user 1's model is used against user 2's data.
table2.add_agreement(
    with_team=tb.config.example_user1["team_id"],
    operation=asset_id,
    session=tb.Session(api_token=tb.config.example_user2["token"], from_default=True),
)


# Use the previously trained XGBoost model to infer against the two
# virtually-joined tables belonging to different organizations.
result = model.infer(
    data=[table1, table2],
    preprocessor=tb.TabularPreprocessor.builder().all_columns(),
    params={"match_column": ["identifier", "id"]},
)
print(result.dataframe)
result.asset.retrieve(save_as="inference_output_asset.zip", overwrite=True)
