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

from pathlib import Path

import tripleblind as tb


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

data_dir = Path.cwd()

# Look for a model Asset ID from a previous run of 1_model_train.py
asset_id = tb.util.load_from("model_asset_id_distributed.out")
model = tb.ModelAsset(asset_id)

test_asset = tb.TableAsset.find(
    "EXAMPLE - Test data - gene regression",
    owned_by=tb.config.example_user2["team_id"],
)

# Perform an inference using the trained model
result = model.infer(
    job_name=f"Regression Inference - {tb.util.timestamp()}",
    data=test_asset,
    params={"security": "smpc"},
)
if result:
    filename = result.asset.retrieve(save_as="linear_model_results.zip", overwrite=True)
    result_df = result.table.dataframe
    print("\nInference results:")
    print("    ", result_df.values.flatten())

else:
    print(f"Inference failed")
