#!/usr/bin/env python3
"""
score_svstr.py — apply the SVSTR-Score confidence model to short-read SV or STR
calls and emit a per-call calibrated confidence score (CS) and tier.

Inference entry point for the released models (sv_model.joblib / str_model.joblib,
each paired with an isotonic calibrator). It loads the trained random forest + its
isotonic calibrator + the feature/config sidecar, then scores a tabular feature
matrix produced by `feature_builder.py` from the caller VCF (short-read VCF +
reference FASTA + static annotation BEDs):

    CS = isotonic_calibrator( RF.predict_proba(X)[:, 1] )    # P(concordant)

TIERS:
  HIGH      CS >= 0.70        (candidate-triage filter)
  MODERATE  0.50 <= CS < 0.70
  WARNING   0.30 <= CS < 0.50
  LOW       CS < 0.30

The score is isotonic-calibrated, so the tier is a pure bucket of the calibrated
CS — there are no heuristic override rules, and STR needs no per-locus catalogue
lookup (its features are self-contained). Missing features (fields a merged or
filtered callset may not carry) are filled with the -99999 sentinel that the
trees were trained to split on.

USAGE
  python score_svstr.py --variant sv  --model-dir . --features sv_features.tsv  --out sv_scored.tsv
  python score_svstr.py --variant str --model-dir . --features str_features.tsv --out str_scored.tsv

Requires the versions in requirements.txt (scikit-learn==1.7.1). Licence: MIT.
"""
import argparse, json, os, sys
import numpy as np
import pandas as pd
import joblib

MISSING = -99999.0
TIER_EDGES = (0.30, 0.50, 0.70)
TIER_NAMES = ("LOW", "WARNING", "MODERATE", "HIGH")


def to_tier(cs):
    return np.asarray(TIER_NAMES)[np.digitize(np.asarray(cs, float), TIER_EDGES, right=False)]


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--variant", choices=["sv", "str"], required=True)
    ap.add_argument("--model-dir", default=".", help="dir with *_config.json + *.joblib")
    ap.add_argument("--features", required=True,
                    help="feature table from feature_builder.py (tsv/csv[.gz])")
    ap.add_argument("--out", required=True)
    ap.add_argument("--raw", action="store_true", help="also emit the uncalibrated RF score (CS_raw)")
    a = ap.parse_args()

    cfg = json.load(open(os.path.join(a.model_dir, f"{a.variant}_config.json")))
    model = joblib.load(os.path.join(a.model_dir, cfg["model_file"]))
    cal = joblib.load(os.path.join(a.model_dir, cfg["calibrator_file"]))
    feats = cfg["features"]

    sep = "\t" if a.features.endswith((".tsv", ".txt", ".gz")) else ","
    df = pd.read_csv(a.features, sep=sep)

    absent = [f for f in feats if f not in df.columns]
    if absent:
        print(f"[warn] {len(absent)} model features absent -> -99999 sentinel: {absent}",
              file=sys.stderr)
        for f in absent:
            df[f] = MISSING

    X = df[feats].astype("float32")
    p = model.predict_proba(X)
    raw = p[:, list(model.classes_).index(1)] if p.shape[1] > 1 else p[:, 0]
    cs = np.clip(cal.predict(raw), 0.0, 1.0)

    out = df.copy()
    if a.raw:
        out["CS_raw"] = np.round(raw, 4)
    out["CS"] = np.round(cs, 4)
    out["tier"] = to_tier(cs)
    out.to_csv(a.out, sep="\t", index=False)

    n = len(out)
    vc = pd.Series(out["tier"]).value_counts()
    hi = int(vc.get("HIGH", 0))
    print(f"{a.variant.upper()}: scored {n:,} calls; HIGH {hi:,} ({hi/n:.1%}) | "
          + " ".join(f"{t}:{int(vc.get(t,0)):,}" for t in TIER_NAMES), file=sys.stderr)
    print(f"wrote {a.out}", file=sys.stderr)


if __name__ == "__main__":
    main()