#!/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_user3["token"], example=True)

# Retrieve some test data to infer against
infer_data = tb.util.download_tripleblind_resource(
    "cali_housing_test.csv",
    save_to_dir="example_data",
    cache_dir="../../.cache",
)

# Look for a model Asset ID from a previous run of 2_train_regression.py
try:
    asset_id = tb.util.load_from("model_asset_id.2b.out")
    model = tb.ModelAsset(asset_id)
    assert model.is_valid
except:
    raise SystemError("You must run 1b_train_regression.py first.")


result = model.infer(
    data=infer_data, params={"security": "smpc", "infer_type": "regression"}
)
if result:
    # Retrieve to review locally
    result.table.retrieve("rf_reg_smpc_infer.zip", overwrite=True)

    print("\nInference results:")
    print(result.table.dataframe.values.flatten())
else:
    print(f"Inference failed")
