"""Train the kcat/Km model.

Stage 1 (no protein language model):
    python train.py --train ../data/train.parquet --out weights
Stage 2 (adds ESM2 protein embeddings; use a GPU in Colab):
    python train.py --train ../data/train.parquet --out weights --esm
Stage 3 (average several models):
    python train.py --train ../data/train.parquet --out weights --esm --seeds 3
"""
import argparse
import json
import os
import time

import lightgbm as lgb
import numpy as np
import pandas as pd
from scipy.stats import spearmanr
from sklearn.model_selection import GroupKFold, GroupShuffleSplit

import features as F

ESM_NAME = "facebook/esm2_t12_35M_UR50D"


def log(*a):
    print(time.strftime("%H:%M:%S"), *a, flush=True)


def te_oof(df, groups, n_splits=5):
    """Substrate-average features for training rows, computed WITHOUT looking at
    the row's own enzyme group (so the model learns to use them honestly)."""
    out = np.full((len(df), 3), np.nan, dtype=np.float32)
    out[:, 0] = 0
    for tr, va in GroupKFold(n_splits).split(df, groups=groups):
        st = F.te_stats(df.iloc[tr])
        out[va] = F.te_apply(df.iloc[va]["smiles"].values, st)
    return out


def make_model(seed, n_estimators, lr):
    return lgb.LGBMRegressor(
        n_estimators=n_estimators, learning_rate=lr, num_leaves=63,
        min_child_samples=20, subsample=0.8, subsample_freq=1,
        colsample_bytree=0.25, reg_lambda=1.0, random_state=seed,
        n_jobs=-1, verbose=-1)


def report(name, y, p):
    scc = spearmanr(y, p).correlation
    rmse = float(np.sqrt(np.mean((y - p) ** 2)))
    log(f"   {name}: SCC={scc:.4f}  RMSE={rmse:.4f}")
    return scc


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train", required=True)
    ap.add_argument("--out", default="weights")
    ap.add_argument("--esm", action="store_true", help="add ESM2 protein embeddings")
    ap.add_argument("--seeds", type=int, default=1)
    ap.add_argument("--lr", type=float, default=0.03)
    ap.add_argument("--max_trees", type=int, default=3000)
    ap.add_argument("--no_holdout", action="store_true", help="skip the local validation run")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)

    df = pd.read_parquet(args.train).reset_index(drop=True)
    log("rows:", len(df), "| unique smiles:", df.smiles.nunique(), "| unique sequences:", df.sequence.nunique())

    usmi = df["smiles"].unique().tolist()
    useq = df["sequence"].unique().tolist()
    log("computing SMILES features ...")
    smi_mat = F.smiles_features(usmi)
    log("computing sequence features ...")
    seq_mat = F.seq_features(useq)

    esm_mat, esm_mean = None, None
    if args.esm:
        esm_dir = os.path.join(args.out, "esm")
        if not os.path.exists(esm_dir):
            from transformers import AutoTokenizer, EsmModel
            log("downloading", ESM_NAME)
            AutoTokenizer.from_pretrained(ESM_NAME).save_pretrained(esm_dir)
            EsmModel.from_pretrained(ESM_NAME).save_pretrained(esm_dir)
        cache = "esm_cache.npz"
        if os.path.exists(cache) and list(np.load(cache, allow_pickle=True)["seqs"]) == useq:
            esm_mat = np.load(cache, allow_pickle=True)["emb"]
            log("loaded ESM embeddings from cache")
        else:
            import torch
            dev = "cuda" if torch.cuda.is_available() else "cpu"
            log(f"computing ESM embeddings on {dev} (this is the slow part) ...")
            esm_mat = F.esm_embed(useq, esm_dir, device=dev,
                                  batch_tokens=16384 if dev == "cuda" else 4096,
                                  n_threads=os.cpu_count() or 2)
            np.savez(cache, seqs=np.array(useq, dtype=object), emb=esm_mat)
        esm_mean = esm_mat.mean(0)
        np.save(os.path.join(args.out, "esm_mean.npy"), esm_mean)

    log("grouping similar sequences (for an honest validation split) ...")
    leader = F.cluster_sequences(useq)
    seq2grp = dict(zip(useq, leader))
    groups = df["sequence"].map(seq2grp).values
    log("sequence groups:", len(set(groups)))

    y_kc = df["log10_kcat"].values
    y_km = df["log10_km"].values
    best_iter = {"kcat": 800, "km": 800}

    # ---------------- local validation run (gives you a trustworthy score) ----
    if not args.no_holdout:
        tr, va = next(GroupShuffleSplit(1, test_size=0.2, random_state=0).split(df, groups=groups))
        dtr, dva = df.iloc[tr].reset_index(drop=True), df.iloc[va].reset_index(drop=True)
        te_tr = te_oof(dtr, groups[tr])
        te_va = F.te_apply(dva["smiles"].values, F.te_stats(dtr))
        Xtr = F.assemble(dtr, usmi, smi_mat, useq, seq_mat, esm_mat, te_tr)
        Xva = F.assemble(dva, usmi, smi_mat, useq, seq_mat, esm_mat, te_va)
        log("holdout features:", Xtr.shape, Xva.shape)
        sccs = []
        for name, col in [("kcat", "log10_kcat"), ("km", "log10_km")]:
            m = make_model(0, args.max_trees, args.lr)
            m.fit(Xtr, dtr[col].values, eval_set=[(Xva, dva[col].values)],
                  callbacks=[lgb.early_stopping(100, verbose=False)])
            best_iter[name] = max(100, int(m.best_iteration_ or args.max_trees))
            sccs.append(report(name, dva[col].values, m.predict(Xva)))
        log(f"LOCAL HOLDOUT mean SCC = {np.mean(sccs):.4f}   (best iters: {best_iter})")

    # ---------------- final training on ALL training rows ---------------------
    log("training final models on all rows ...")
    te_all = te_oof(df, groups)
    stats = F.te_stats(df)
    stats.to_parquet(os.path.join(args.out, "te_stats.parquet"))
    X = F.assemble(df, usmi, smi_mat, useq, seq_mat, esm_mat, te_all)
    for name, y in [("kcat", y_kc), ("km", y_km)]:
        for s in range(args.seeds):
            m = make_model(s, int(best_iter[name] * 1.15), args.lr)
            m.fit(X, y)
            m.booster_.save_model(os.path.join(args.out, f"{name}_{s}.txt"))
            log(f"saved {name}_{s}.txt")

    meta = {"use_esm": bool(args.esm), "n_seeds": args.seeds,
            "mean_kcat": float(y_kc.mean()), "mean_km": float(y_km.mean())}
    with open(os.path.join(args.out, "meta.json"), "w") as f:
        json.dump(meta, f)
    log("DONE. Weights are in", args.out)


if __name__ == "__main__":
    main()