#!/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()


def print_in_columns(lst, indent=0):
    import shutil

    int_width = 0
    dec_width = 0
    for v in lst:
        w = len(str(v))
        iw = len(str(int(v)))
        int_width = max(int_width, iw)
        dec_width = max(dec_width, w - iw - 1)
    max_width = int_width + dec_width + 1

    term_width = shutil.get_terminal_size().columns
    cols = (term_width - indent - 2) // (max_width + 1)

    if (max_width + 1) * len(lst) + indent + 4 < term_width:
        print(f"{' '*indent}[ ", end="")
        for v in lst:
            print(f"{v} ", end="")
        print("]")
    else:
        col = 0
        print(f"{' '*indent}[")
        for v in lst:
            if col == 0:
                print(" " * (indent + 1), end="")
            # print(f"0:>{max_width}".format(str(v)), end="")
            padded = f"{v:.{dec_width}f}"
            print(f"{padded:>{max_width}} ", end="")
            col += 1
            if col >= cols:
                print()
                col = 0
        print(f"\n{' '*indent}]")


# retrieve the dataset using tripleblind to use for local inference
test_asset = tb.TableAsset.find(
    "EXAMPLE - Test data - gene regression",
    owned_by=tb.config.example_user2["team_id"],
)
test_asset.retrieve(save_as="test_asset.zip", overwrite=True)
df = tb.Package.load("test_asset.zip").records()

pack = tb.Package.load("gene_model.zip")
model = pack.model()

# Perform an inference using the trained model
print("Inference result:")
results = model.predict(df.values)
print_in_columns(results, indent=4)

print("Coefficients:")
print(model.coef_)
