"""Inference entry point:  python /app/predict.py --input X.parquet --output Y.csv"""
import argparse
import json
import os
import time

import numpy as np
import pandas as pd

import features as F

HERE = os.path.dirname(os.path.abspath(__file__))
W = os.path.join(HERE, "weights")
N_THREADS = int(os.environ.get("NTHREADS", 2))   # the evaluator gives us 2 vCPUs
ESM_BUDGET_S = 480                                # max 8 minutes for ESM, then fall back


def main():
    t0 = time.time()
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--output", required=True)
    args = ap.parse_args()

    df = pd.read_parquet(args.input).reset_index(drop=True)
    df["smiles"] = df["smiles"].fillna("").astype(str)
    df["sequence"] = df["sequence"].fillna("").astype(str)
    meta = json.load(open(os.path.join(W, "meta.json")))

    usmi = df["smiles"].unique().tolist()
    useq = df["sequence"].unique().tolist()
    smi_mat = F.smiles_features(usmi)
    seq_mat = F.seq_features(useq)
    print(f"[predict] features done {time.time() - t0:.0f}s", flush=True)

    esm_mat = None
    if meta["use_esm"]:
        fb = np.load(os.path.join(W, "esm_mean.npy"))
        esm_mat = F.esm_embed(useq, os.path.join(W, "esm"), fallback=fb,
                              budget_s=ESM_BUDGET_S, n_threads=N_THREADS)
        print(f"[predict] esm done {time.time() - t0:.0f}s", flush=True)

    stats = pd.read_parquet(os.path.join(W, "te_stats.parquet"))
    ref_bits = np.load(os.path.join(W, "ref_bits.npy")).astype(np.float32)
    smi_index = {s: i for i, s in enumerate(usmi)}
    bits_u = (smi_mat[:, :F.N_MORGAN] > 0).astype(np.float32)
    te = F.te_full(df["smiles"].values, smi_index, bits_u, stats, ref_bits)
    X = F.assemble(df, usmi, smi_mat, useq, seq_mat, esm_mat, te)

    import lightgbm as lgb
    out = pd.DataFrame({"id": df["id"]})
    for name, col, mean_key in [("kcat", "log10_kcat", "mean_kcat"), ("km", "log10_km", "mean_km")]:
        preds = []
        for s in range(meta["n_seeds"]):
            b = lgb.Booster(model_file=os.path.join(W, f"{name}_{s}.txt"))
            preds.append(b.predict(X, num_threads=N_THREADS))
        p = np.mean(preds, axis=0)
        p = np.nan_to_num(p, nan=meta[mean_key], posinf=meta[mean_key], neginf=meta[mean_key])
        out[col] = np.clip(p, -50, 50)

    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
    out.to_csv(args.output, index=False)
    print(f"[predict] wrote {len(out)} rows in {time.time() - t0:.0f}s", flush=True)


if __name__ == "__main__":
    main()