"""5+ 술식 stratification: RYGB / SG / OAGB / SADI / DJB.

Per-procedure:
  - N, mean age/BMI/T2DM 비율
  - Procedure-specific complication rates
    * RYGB : leak, marginal_ulcer, dumping, delayed hypoglycemia (POD30+)
    * SG   : staple leak, stricture, GERD, delayed hypoglycemia
    * OAGB : leak, bile reflux
    * SADI : leak, malabsorption
    * DJB  : weight loss-vs-rest comparison
  - 30-day readmission / mortality rate
  - NSQIP/MBSAQIP-analog risk-adjusted Observed / Expected (O/E) ratio
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional

from .ingest import Patient, POD03Row, POD430Row


# Naive expected (E) rates per procedure — published-literature anchor (synthetic)
EXPECTED_RATES = {
    "RYGB": {
        "leak": 0.014, "marginal_ulcer": 0.040, "dumping_any": 0.20,
        "readmit": 0.062, "mortality": 0.0025,
    },
    "SG": {
        "leak": 0.020, "stricture": 0.020, "GERD_progress": 0.10,
        "readmit": 0.045, "mortality": 0.0015,
    },
    "OAGB": {
        "leak": 0.014, "bile_reflux": 0.04,
        "readmit": 0.055, "mortality": 0.0020,
    },
    "SADI": {
        "leak": 0.020, "malabsorption": 0.10,
        "readmit": 0.060, "mortality": 0.0025,
    },
    "DJB": {
        "leak": 0.012, "readmit": 0.050, "mortality": 0.0020,
    },
}


@dataclass
class ProcedureStrat:
    procedure: str
    n: int
    mean_age: float
    mean_bmi: float
    pct_female: float
    pct_t2dm: float
    leak_rate_pct: float
    marginal_ulcer_pct: float
    stricture_late_pct: float
    dumping_any_pct: float
    bile_reflux_pct: float
    malabsorption_pct: float
    los_mean_d: float
    readmit_30d_pct: float
    mortality_30d_pct: float
    oe_leak: Optional[float]
    oe_readmit: Optional[float]
    oe_mortality: Optional[float]


def _safe(num: int, denom: int) -> float:
    return round(100.0 * num / denom, 2) if denom else 0.0


def _safe_oe(observed_rate: float, expected_rate: float) -> Optional[float]:
    if expected_rate <= 0:
        return None
    return round(observed_rate / (expected_rate * 100.0), 2)


def stratify_by_procedure(patients: List[Patient],
                          pod03: List[POD03Row],
                          pod430: List[POD430Row]) -> List[ProcedureStrat]:
    by_pid = {p.patient_id: p for p in patients}
    pod03_map = {r.patient_id: r for r in pod03}
    pod430_map = {r.patient_id: r for r in pod430}

    grp: Dict[str, List[Patient]] = {}
    for p in patients:
        grp.setdefault(p.procedure, []).append(p)

    out: List[ProcedureStrat] = []
    for proc in sorted(grp):
        plist = grp[proc]
        n = len(plist)
        ages = [p.age for p in plist]
        bmis = [p.bmi_pre for p in plist]
        n_f = sum(1 for p in plist if p.sex == "F")
        n_dm = sum(1 for p in plist if p.t2dm)
        los_vals = [p.los_days for p in plist]

        n_leak = sum(1 for p in plist
                     if (pod03_map.get(p.patient_id)
                         and (pod03_map[p.patient_id].leak
                              or pod03_map[p.patient_id].staple_leak)))
        n_marg = sum(1 for p in plist
                     if pod430_map.get(p.patient_id)
                     and pod430_map[p.patient_id].marginal_ulcer)
        n_str = sum(1 for p in plist
                    if pod430_map.get(p.patient_id)
                    and pod430_map[p.patient_id].stricture_late)
        n_dump = sum(1 for p in plist
                     if pod430_map.get(p.patient_id) and
                     (pod430_map[p.patient_id].dumping_early
                      or pod430_map[p.patient_id].dumping_late))
        n_bile = sum(1 for p in plist
                     if pod430_map.get(p.patient_id)
                     and pod430_map[p.patient_id].bile_reflux)
        n_mal = sum(1 for p in plist
                    if pod430_map.get(p.patient_id)
                    and pod430_map[p.patient_id].malabsorption)
        n_readmit = sum(1 for p in plist
                        if pod430_map.get(p.patient_id)
                        and pod430_map[p.patient_id].readmit_30d)
        n_mort = sum(1 for p in plist if p.died_30d)

        leak_rate = n_leak / n if n else 0.0
        readmit_rate = n_readmit / n if n else 0.0
        mort_rate = n_mort / n if n else 0.0

        exp = EXPECTED_RATES.get(proc, {})
        oe_leak = _safe_oe(leak_rate * 100.0,
                           exp.get("leak", 0.0) * 100.0) if exp else None
        oe_rd = _safe_oe(readmit_rate * 100.0,
                         exp.get("readmit", 0.0) * 100.0) if exp else None
        oe_mt = _safe_oe(mort_rate * 100.0,
                         exp.get("mortality", 0.0) * 100.0) if exp else None

        out.append(ProcedureStrat(
            procedure=proc,
            n=n,
            mean_age=round(sum(ages) / n, 1) if n else 0.0,
            mean_bmi=round(sum(bmis) / n, 1) if n else 0.0,
            pct_female=_safe(n_f, n),
            pct_t2dm=_safe(n_dm, n),
            leak_rate_pct=_safe(n_leak, n),
            marginal_ulcer_pct=_safe(n_marg, n),
            stricture_late_pct=_safe(n_str, n),
            dumping_any_pct=_safe(n_dump, n),
            bile_reflux_pct=_safe(n_bile, n),
            malabsorption_pct=_safe(n_mal, n),
            los_mean_d=round(sum(los_vals) / n, 1) if n else 0.0,
            readmit_30d_pct=_safe(n_readmit, n),
            mortality_30d_pct=_safe(n_mort, n),
            oe_leak=oe_leak,
            oe_readmit=oe_rd,
            oe_mortality=oe_mt,
        ))
    return out