#!/usr/bin/env python3
"""Standalone inference for the Inflammatory Aging Clock TabM student models.

Runs the distilled TabM student with public dependencies only:

    pip install torch tabm rtdl_num_embeddings numpy pandas

Usage
-----
    python predict.py --input proteins.csv --output ages.csv

Input
-----
A CSV/TSV with one row per sample and one column per protein. Column names must
match the model's feature names (e.g. ``OID00471_CXCL8`` for MARS, the SomaScan
SeqId-style names for ROSMAP — see ``<model_dir>/meta.json`` -> ``feature_name``).
An optional first ``sample_id`` (or ``SAMPLE.ID``) column is carried through to
the output. Missing proteins are an error; NaN cells are median-imputed for
models trained with a median-impute preprocessor (MARS), and rejected otherwise.

Output
------
A CSV with columns ``sample_id, predicted_age``.
"""
import argparse
import json
import os
import sys

import numpy as np
import pandas as pd
import torch


# --- pinball5 reparameterization (numpy mirror of the training code) ---------
PINBALL5_ALPHAS = np.array([0.1, 0.25, 0.5, 0.75, 0.9], dtype=np.float64)


def _reparameterize_pinball5_np(raw_5):
    """raw (..., 5) -> monotone quantiles (q01, q025, q05, q075, q09)."""
    def softplus(x):
        return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)
    m = raw_5[..., 0]
    s1 = softplus(raw_5[..., 1]); s2 = softplus(raw_5[..., 2])
    s3 = softplus(raw_5[..., 3]); s4 = softplus(raw_5[..., 4])
    return np.stack([m - s1 - s2, m - s2, m, m + s3, m + s3 + s4], axis=-1)


def _pinball5_point_estimate_np(raw_5, alphas=PINBALL5_ALPHAS):
    """Trapezoidal-mean point estimate from raw (..., 5). Matches the teacher's mean."""
    q = _reparameterize_pinball5_np(raw_5)
    slope_low = (q[..., 1] - q[..., 0]) / (alphas[1] - alphas[0])
    slope_high = (q[..., 4] - q[..., 3]) / (alphas[4] - alphas[3])
    q0 = q[..., 0] - slope_low * alphas[0]
    q1 = q[..., 4] + slope_high * (1.0 - alphas[4])
    q_full = np.concatenate([q0[..., None], q, q1[..., None]], axis=-1)
    a_full = np.array([0.0, *alphas, 1.0])
    trapz = np.trapezoid if hasattr(np, "trapezoid") else np.trapz  # numpy>=2 vs <2
    return trapz(q_full, a_full, axis=-1).astype(np.float32, copy=False)


# --- preprocessing (median-impute + observed-mask channels) ------------------
def _apply_preproc(X, npz_path):
    """Median-impute NaN and append a mask channel for partially observed columns."""
    d = np.load(npz_path)
    medians = d["medians"].astype(np.float32)
    has_nan_cols = (d["has_nan_cols"].astype(bool) if "has_nan_cols" in d.files
                    else np.ones(medians.shape[0], dtype=bool))
    X = np.asarray(X, dtype=np.float32)
    nan = np.isnan(X)
    X_imp = np.where(nan, medians, X).astype(np.float32, copy=False)
    if not has_nan_cols.any():
        return X_imp
    M = (~nan[:, has_nan_cols]).astype(np.float32)
    return np.concatenate([X_imp, M], axis=1)


# --- model construction ------------------------------------------------------
def _build_model(meta, weights_path):
    import tabm as _tabm
    state = torch.load(weights_path, map_location="cpu", weights_only=False)
    if isinstance(state, dict) and "state_dict" in state:
        sd, saved_bins = state["state_dict"], state.get("bins")
    else:
        sd, saved_bins = state, None
    arch = meta["arch"]
    em = arch["embed_meta"]
    if em.get("kind") == "piecewise_linear":
        import rtdl_num_embeddings as _rne
        if saved_bins is None:
            raise ValueError("Model declares piecewise_linear embeddings but .pt has no 'bins'.")
        emb = _rne.PiecewiseLinearEmbeddings(
            saved_bins, d_embedding=int(em["d_embedding"]),
            activation=bool(em.get("activation", False)), version=em.get("version", "B"),
        )
    elif em.get("kind") in (None, "none"):
        emb = None
    else:
        raise ValueError(f"Unsupported embed_meta.kind={em.get('kind')!r}")
    model = _tabm.TabM(
        n_num_features=int(arch["d_in"]), cat_cardinalities=[],
        d_out=int(arch["n_outputs"]), num_embeddings=emb, **arch["tabm_kwargs"],
    )
    model.load_state_dict(sd)
    model.eval()
    return model


def _predict_scaled(model, X, distill_loss, batch_size=512):
    """Forward TabM -> (n,) scaled-y prediction, averaged over the k ensemble members."""
    out_chunks = []
    with torch.no_grad():
        for i in range(0, len(X), batch_size):
            xb = torch.from_numpy(X[i:i + batch_size].astype(np.float32, copy=False))
            out = model(x_num=xb, x_cat=None).float().cpu().numpy()  # (B, k, d_out)
            if distill_loss == "pinball5":
                point_bk = _pinball5_point_estimate_np(out)          # (B, k)
                out_chunks.append(point_bk.mean(axis=1))
            else:
                out_chunks.append(out.squeeze(-1).mean(axis=1))
    return np.concatenate(out_chunks, axis=0)


# --- I/O ---------------------------------------------------------------------
def _read_table(path):
    sep = "\t" if path.lower().endswith((".tsv", ".txt", ".tsv.gz", ".txt.gz")) else ","
    return pd.read_csv(path, sep=sep)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model_dir", default=".",
                    help="Folder containing meta.json + models/ (default: current directory).")
    ap.add_argument("--input", required=True, help="CSV/TSV of proteins (samples x features).")
    ap.add_argument("--output", required=True, help="Where to write predictions CSV.")
    args = ap.parse_args()

    meta = json.load(open(os.path.join(args.model_dir, "meta.json")))
    feats = meta["feature_name"]

    df = _read_table(args.input)
    # Carry a sample id column through if present.
    id_col = next((c for c in ("sample_id", "SAMPLE.ID", "SampleID", "ID") if c in df.columns), None)
    sample_ids = df[id_col].astype(str).values if id_col else np.arange(len(df)).astype(str)

    # The model needs EVERY trained protein present, by exact name. Column order
    # does not matter (selected/reordered below). Hard-fail on any missing one.
    feat_set = set(feats)
    missing = [f for f in feats if f not in df.columns]
    if missing:
        sys.exit(f"ERROR: input is missing {len(missing)}/{len(feats)} required protein "
                 f"column(s), e.g. {missing[:5]}. Column names must match "
                 f"{args.model_dir}/meta.json -> feature_name exactly.")
    # Warn about input columns the model does not use (often a naming mismatch).
    extra = [c for c in df.columns if c not in feat_set and c != id_col]
    if extra:
        print(f"WARNING: ignoring {len(extra)} input column(s) not used by the model, "
              f"e.g. {extra[:5]}.", file=sys.stderr)
    print(f"Matched all {len(feats)} required feature columns.", file=sys.stderr)
    X = df[feats].to_numpy(dtype=np.float32)  # (n, n_features_raw), column order = feats

    if meta.get("preproc"):
        X = _apply_preproc(X, os.path.join(args.model_dir, meta["preproc"]["file"]))
    elif np.isnan(X).any():
        sys.exit("ERROR: input contains NaN but this model has no imputer. "
                 "Provide complete protein values.")

    model = _build_model(meta, os.path.join(args.model_dir, meta["student_weights"]))
    scaled = _predict_scaled(model, X, meta["regression_distill_loss"])
    age = scaled * meta["y_scaler"]["scale"] + meta["y_scaler"]["mean"]

    out = pd.DataFrame({"sample_id": sample_ids, "predicted_age": age})
    out.to_csv(args.output, index=False)
    print(f"Wrote {len(out)} predictions to {args.output}")


if __name__ == "__main__":
    main()