"""
runtime.py
==========
Web 서비스용 런타임 래퍼. webtool_runtime.py 를 HF Space 에 맞게 이식했다.

로드 대상 (기동 시 1회):
    webtool_core.py          모델 정의 + 상수 + feature + ensemble
    webtool_baseline.json    동결된 risk -> S(H) 매핑 (+ bootstrap CI 계수)
    weights/seed_*/fold_*.pt 31 seed x 10 fold 앙상블 (310 checkpoints)

입력  : T stage(1..6), station별 전이 count, station별 harvest(절제 node 수)
출력  : risk_mean, risk_std, SCR_3, entropy, total_meta,
        surv_60m = 기대 5년 생존확률(+95% CI), extrapolated flag
"""
from __future__ import annotations

import json
import os
import threading
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import torch

from .webtool_core import (
    STATIONS, T_MAP, load_canonical,
    build_ln_features, build_t_onehot, scr_and_entropy, MultiSeedEnsemble,
)


def _n_label(total_meta: float) -> str:
    """전이 node 총수 -> AJCC N stage (N0/N1/N2/N3a/N3b)."""
    n = int(round(total_meta or 0))
    if n == 0:
        return "N0"
    if n <= 2:
        return "N1"
    if n <= 6:
        return "N2"
    if n <= 15:
        return "N3a"
    return "N3b"

# --------------------------------------------------------------------------- #
# 경로 / 설정
# --------------------------------------------------------------------------- #
BUNDLE_ROOT = Path(__file__).resolve().parents[1]

CONFIG = {
    "WEIGHTS_ROOT": os.environ.get("WEIGHTS_ROOT", str(BUNDLE_ROOT / "weights")),
    "CANON_SEED": int(os.environ.get("CANON_SEED", "9999")),
    "SEEDS": None,  # None 이면 seed_* 전부 자동 탐색 (baseline 동결 시점과 동일해야 함)
    "BATCH": int(os.environ.get("BATCH", "256")),
    "DEVICE": "cuda" if torch.cuda.is_available() else "cpu",
    "ARTIFACT_PATH": os.environ.get(
        "ARTIFACT_PATH", str(BUNDLE_ROOT / "webtool_baseline.json")
    ),
}
DEVICE = torch.device(CONFIG["DEVICE"])

# T stage 코드(1..6) -> 라벨. UI/스키마에서 사용.
T_OPTIONS = [{"value": k, "label": v} for k, v in sorted(T_MAP.items())]

# 지연 초기화 상태 (thread-safe)
_STATE: dict[str, Any] = {}
_INIT_LOCK = threading.Lock()


def _ensure_baseline() -> dict[str, Any]:
    """baseline json + canonical config 만 로드한다(저렴). 앙상블은 로드하지 않음.
    schema/health 가 앙상블 로딩을 기다리지 않도록 분리한다."""
    if "art" in _STATE:
        return _STATE
    with _INIT_LOCK:
        if "art" in _STATE:
            return _STATE
        print("[runtime] loading baseline + canonical config …", flush=True)
        with open(CONFIG["ARTIFACT_PATH"]) as f:
            art = json.load(f)
        model_config, t_cols = load_canonical(CONFIG["WEIGHTS_ROOT"], CONFIG["CANON_SEED"])
        # 동결 시점과 런타임의 구조 및 T 인코딩 정합을 강제한다.
        assert art["t_cols"] == t_cols, "runtime t_cols 가 baseline 과 불일치"
        assert art["model_config"] == model_config, "runtime model_config 가 baseline 과 불일치"
        _STATE.update(
            art=art, model_config=model_config, t_cols=t_cols,
            horizon=art["horizon_months"],
        )
        print(f"[runtime] baseline OK. horizon={art['horizon_months']}mo", flush=True)
    return _STATE


def init() -> dict[str, Any]:
    """전체 앙상블(310 checkpoints)까지 로드한다. idempotent, thread-safe.
    기동 지연을 피하려고 서버는 이 함수를 백그라운드 스레드에서 호출한다."""
    if _STATE.get("ready"):
        return _STATE
    _ensure_baseline()
    with _INIT_LOCK:
        if _STATE.get("ready"):
            return _STATE
        print("[runtime] building ensemble (310 checkpoints) …", flush=True)
        ensemble = MultiSeedEnsemble(
            weights_root=CONFIG["WEIGHTS_ROOT"],
            model_config=_STATE["model_config"],
            batch=CONFIG["BATCH"],
            device=DEVICE,
            seeds=CONFIG["SEEDS"],
        )
        _STATE.update(ensemble=ensemble, ready=True, n_seeds=len(ensemble.seeds))
        print(f"[runtime] ensemble ready: {len(ensemble.seeds)} seeds", flush=True)
    return _STATE


# --------------------------------------------------------------------------- #
# risk -> 기대 survival (+ CI). webtool_runtime.py 수식 그대로.
# --------------------------------------------------------------------------- #
def _expected_survival(risk_series: pd.Series, art: dict) -> tuple[pd.Series, pd.Series]:
    r = risk_series.astype(float)
    s = art["S0_H"] ** np.exp(art["beta"] * (r - art["risk_center"]))
    oor = (r < art["risk_ref_min"]) | (r > art["risk_ref_max"])
    return (
        pd.Series(s, index=r.index, name=f"surv_{art['horizon_months']}m"),
        pd.Series(oor, index=r.index, name="extrapolated"),
    )


def _expected_survival_ci(
    risk_series: pd.Series, per_seed: pd.DataFrame | None, art: dict, mode: str = "cox"
) -> pd.DataFrame:
    r = risk_series.astype(float)
    boot = art["boot"]
    B, N = len(boot), len(r)
    M = np.empty((B, N))

    seeds_arr = None
    if mode == "full":
        assert per_seed is not None, "mode='full' 에는 per_seed DataFrame 필요"
        seeds_arr = per_seed.reindex(r.index).values  # [N, n_seeds]

    for j, bp in enumerate(boot):
        if mode == "full":
            rng = np.random.default_rng(1000 + j)
            n_seed = seeds_arr.shape[1]
            r_eval = np.array(
                [
                    rng.choice(seeds_arr[i], size=n_seed, replace=True).mean()
                    for i in range(N)
                ]
            )
        else:
            r_eval = r.values
        M[j] = bp["S0_H"] ** np.exp(bp["beta"] * (r_eval - bp["risk_center"]))

    a = art["ci_alpha"]
    lo = np.percentile(M, 100 * a / 2, axis=0)
    hi = np.percentile(M, 100 * (1 - a / 2), axis=0)
    return pd.DataFrame({"lo": lo, "hi": hi}, index=r.index)


# --------------------------------------------------------------------------- #
# End-to-end 채점 (DataFrame 반환)
# --------------------------------------------------------------------------- #
def score_frame(
    tstage: pd.Series,
    count_df: pd.DataFrame,
    harvest_df: pd.DataFrame,
    ci: bool = True,
    ci_mode: str = "cox",
) -> pd.DataFrame:
    """다수 환자 채점. index 보존. 컬럼 순서/의미는 webtool_runtime 과 동일."""
    st = init()
    art, ensemble, t_cols = st["art"], st["ensemble"], st["t_cols"]

    ln = build_ln_features(count_df, harvest_df)
    t_df = build_t_onehot(tstage.reindex(count_df.index), t_cols)
    risk, per_seed = ensemble.predict(ln, t_df, return_per_seed=True)

    se = scr_and_entropy(count_df)
    s_H, oor = _expected_survival(risk, art)
    H = art["horizon_months"]

    out = pd.DataFrame(index=count_df.index)
    out["Tstage"] = tstage.reindex(count_df.index).astype(int)
    out["total_meta"] = se["total_meta"]
    out["risk_mean"] = risk
    out["risk_std"] = per_seed.std(axis=1)
    out["SCR_3"] = se["SCR_3"]
    out["entropy"] = se["entropy"]
    out[f"surv_{H}m"] = s_H
    if ci:
        cid = _expected_survival_ci(risk, per_seed=per_seed, art=art, mode=ci_mode)
        out[f"surv_{H}m_lo"] = cid["lo"]
        out[f"surv_{H}m_hi"] = cid["hi"]
        out["ci_mode"] = ci_mode
    out["extrapolated"] = oor
    return out


# --------------------------------------------------------------------------- #
# 입력 정규화 헬퍼 (dict 형태 -> STATIONS 순서 DataFrame)
# --------------------------------------------------------------------------- #
def _station_df(mapping: dict[str, float], index) -> pd.DataFrame:
    """{station: 값} -> [1,16] DataFrame. 미기재 station 은 0."""
    mapping = mapping or {}
    return pd.DataFrame(
        [[float(mapping.get(s, 0) or 0) for s in STATIONS]],
        columns=STATIONS,
        index=index,
    )


def _clean(v):
    """numpy/pandas 스칼라 -> JSON 직렬화 가능한 파이썬 값. NaN -> None."""
    if v is None:
        return None
    if isinstance(v, (np.floating, float)):
        v = float(v)
        return None if np.isnan(v) else v
    if isinstance(v, (np.integer,)):
        return int(v)
    if isinstance(v, (np.bool_, bool)):
        return bool(v)
    return v


def _row_to_result(row: pd.Series, patient_id: str, H: int) -> dict[str, Any]:
    """score_frame 한 행 -> API 응답 dict."""
    return {
        "id": patient_id,
        "Tstage": _clean(row["Tstage"]),
        "T_label": T_MAP.get(int(row["Tstage"])),
        "total_meta": _clean(row["total_meta"]),
        "risk_mean": _clean(row["risk_mean"]),
        "risk_std": _clean(row["risk_std"]),
        "SCR_3": _clean(row["SCR_3"]),
        "entropy": _clean(row["entropy"]),
        "surv": _clean(row[f"surv_{H}m"]),
        "surv_lo": _clean(row.get(f"surv_{H}m_lo")),
        "surv_hi": _clean(row.get(f"surv_{H}m_hi")),
        "extrapolated": _clean(row["extrapolated"]),
        "horizon_months": H,
    }


def patient_curve(risk: float) -> dict[str, Any] | None:
    """환자 위험점수 -> 시간축 생존곡선 S(t|risk) + 95% CI band + 참조 KM.
    S(t|risk) = S_center(t) ** exp(beta*(risk - risk_center)).  S_center(60)=S0_H 이므로
    곡선의 60개월 값은 point surv 와 정확히 일치한다. baseline 아티팩트가 curve 를
    포함하지 않으면 None."""
    art = init()["art"]
    if "curve_base" not in art:
        return None
    beta, rc = art["beta"], art["risk_center"]
    t = np.asarray(art["curve_months"], dtype=int)
    base = np.asarray(art["curve_base"], dtype=float)
    s = base ** np.exp(beta * (risk - rc))

    bb = np.asarray(art["curve_boot_base"], dtype=float)          # [B, T]
    bbeta = np.asarray(art["curve_boot_beta"], dtype=float)[:, None]
    bcen = np.asarray(art["curve_boot_center"], dtype=float)[:, None]
    sb = bb ** np.exp(bbeta * (risk - bcen))                      # [B, T]
    a = art.get("ci_alpha", 0.05)
    lo = np.percentile(sb, 100 * a / 2, axis=0)
    hi = np.percentile(sb, 100 * (1 - a / 2), axis=0)

    # 생존함수는 수술 시점(t=0)에서 정의상 1. baseline curve 의 0개월 값이 정확히
    # 1 이 아닐 수 있으므로 그래프와 RMST 적분 모두에서 강제로 보정한다.
    s[0] = 1.0
    lo[0] = 1.0
    hi[0] = 1.0
    sb[:, 0] = 1.0

    # RMST_H = ∫_0^H S(t) dt (trapezoidal, 월 단위). 항상 0..H 범위.
    # CI 는 최종 lo/hi 곡선을 적분하는 게 아니라, bootstrap replicate 곡선별
    # RMST 분포의 2.5/97.5 percentile 을 쓴다.
    tf = t.astype(float)
    _trapz = getattr(np, "trapezoid", np.trapz)
    rmst = float(_trapz(s, tf))
    rmst_boot = _trapz(sb, tf, axis=1)
    rmst_lo, rmst_hi = (
        float(x)
        for x in np.percentile(rmst_boot, [100 * a / 2, 100 * (1 - a / 2)])
    )

    r5 = lambda arr: [round(float(x), 5) for x in arr]
    return {
        "t": t.tolist(),
        "s": r5(s),
        "lo": r5(lo),
        "hi": r5(hi),
        "ref_km": r5(art["ref_km"]) if "ref_km" in art else None,
        "rmst": round(rmst, 2),
        "rmst_lo": round(rmst_lo, 2),
        "rmst_hi": round(rmst_hi, 2),
    }


def score_single(
    tstage: int,
    counts: dict[str, float],
    harvests: dict[str, float],
    patient_id: str = "patient",
    ci_mode: str = "cox",
) -> dict[str, Any]:
    """단일 환자 채점 -> API 응답 dict (+ 시간축 생존곡선)."""
    st = init()
    idx = [patient_id]
    count_df = _station_df(counts, idx)
    harvest_df = _station_df(harvests, idx)
    t_series = pd.Series([int(tstage)], index=idx)
    out = score_frame(t_series, count_df, harvest_df, ci=True, ci_mode=ci_mode)
    res = _row_to_result(out.iloc[0], patient_id, st["horizon"])
    res["curve"] = patient_curve(float(res["risk_mean"]))

    # RMST (H-개월 restricted mean survival time) 를 상위 레벨로 승격.
    # 고정 키(rmst*)는 프론트/일반 소비자용, rmst_{H}m* 는 horizon 명시용.
    H = st["horizon"]
    cur = res.get("curve")
    if cur and cur.get("rmst") is not None:
        res["rmst"] = cur["rmst"]
        res["rmst_lo"] = cur["rmst_lo"]
        res["rmst_hi"] = cur["rmst_hi"]
        res[f"rmst_{H}m"] = cur["rmst"]
        res[f"rmst_{H}m_lo"] = cur["rmst_lo"]
        res[f"rmst_{H}m_hi"] = cur["rmst_hi"]
        res["rmst_horizon_months"] = H

    art = st["art"]
    # TNM 분류 (T + total_meta 로부터 N 유도 -> TNM 그룹/stage)
    t_clin = T_MAP.get(int(res["Tstage"]))
    t_grp = art.get("tmap_tnm", {}).get(str(int(res["Tstage"])), t_clin)
    n_lab = _n_label(res["total_meta"])
    group = art.get("tnm_dict", {}).get(f"{t_grp}|{n_lab}")
    stage = art.get("tnm_label_map", {}).get(str(group)) if group else None
    res["tnm"] = {"t": t_clin, "n": n_lab, "tn": f"{t_clin}{n_lab}",
                  "group": group, "stage": stage}

    # SCR: 절대 기준 0.75 (high/low). entropy: ln2/ln4 로 low/intermediate/high.
    import math
    scr, ent = res["SCR_3"], res["entropy"]
    scr_thr = 0.75
    res["scr_flag"] = None if scr is None else ("high" if scr >= scr_thr else "low")
    res["scr_threshold"] = scr_thr
    if ent is None:
        res["entropy_level"] = None
    elif ent < math.log(2):
        res["entropy_level"] = "low"
    elif ent < math.log(4):
        res["entropy_level"] = "intermediate"
    else:
        res["entropy_level"] = "high"
    res["entropy_thresholds"] = [round(math.log(2), 4), round(math.log(4), 4)]
    return res


# --------------------------------------------------------------------------- #
# CSV 배치 채점
# --------------------------------------------------------------------------- #
def csv_columns() -> list[str]:
    """배치 CSV 표준 컬럼: id, Tstage, meta_<station>, harv_<station> x 16."""
    cols = ["id", "Tstage"]
    for s in STATIONS:
        cols += [f"meta_{s}", f"harv_{s}"]
    return cols


def score_csv(df: pd.DataFrame, ci_mode: str = "cox") -> list[dict[str, Any]]:
    """표준 CSV DataFrame 채점. 누락 station 컬럼은 0 으로 채운다."""
    st = init()
    H = st["horizon"]
    df = df.copy()

    if "Tstage" not in df.columns:
        raise ValueError("CSV 에 'Tstage' 컬럼이 필요합니다 (값 1..6).")
    if "id" not in df.columns:
        df["id"] = [f"row_{i:04d}" for i in range(len(df))]
    ids = df["id"].astype(str).tolist()
    index = pd.RangeIndex(len(df))

    count_df = pd.DataFrame(index=index, columns=STATIONS, dtype=float)
    harvest_df = pd.DataFrame(index=index, columns=STATIONS, dtype=float)
    for s in STATIONS:
        count_df[s] = pd.to_numeric(df.get(f"meta_{s}", 0), errors="coerce").fillna(0).values
        harvest_df[s] = pd.to_numeric(df.get(f"harv_{s}", 0), errors="coerce").fillna(0).values

    t_series = pd.Series(
        pd.to_numeric(df["Tstage"], errors="coerce").fillna(1).astype(int).values,
        index=index,
    )
    out = score_frame(t_series, count_df, harvest_df, ci=True, ci_mode=ci_mode)
    return [_row_to_result(out.iloc[i], ids[i], H) for i in range(len(out))]