"""Inspecciones sanitarias de restaurantes de NYC — priorización de inspecciones.

Pregunta de negocio: ¿qué restaurantes tienen más riesgo de suspender (SCORE >= 14,
es decir, no sacar una A) en su próxima inspección inicial de ciclo?

Solo se usan datos conocidos ANTES de la inspección (historial del restaurante,
cocina, distrito). SCORE y GRADE de la propia inspección quedan fuera: GRADE es una
función directa de SCORE y usarlos daba un ROC-AUC de 0,9998 que no servía de nada.

Uso:
    python analysis/nyc_inspections.py raw/DOHMH_New_York_City_Restaurant_Inspection_Results.csv
"""
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import roc_auc_score

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "data" / "nyc.json"
SPLIT_DATE = "2025-07-01"
FAIL_SCORE = 14  # a partir de 14 puntos ya no es A (normativa DOHMH)
INITIAL = "Cycle Inspection / Initial Inspection"


def load(path):
    cols = ["CAMIS", "BORO", "CUISINE DESCRIPTION", "INSPECTION DATE", "ACTION", "VIOLATION CODE",
            "VIOLATION DESCRIPTION", "CRITICAL FLAG", "SCORE", "GRADE", "INSPECTION TYPE"]
    df = pd.read_csv(path, usecols=cols, dtype=str)
    df["date"] = pd.to_datetime(df["INSPECTION DATE"], format="%m/%d/%Y", errors="coerce")
    # 01/01/1900 = restaurante aún sin inspeccionar
    df = df[(df["date"].dt.year >= 2022) & df["INSPECTION TYPE"].notna() & (df["BORO"] != "0")]
    df["SCORE"] = pd.to_numeric(df["SCORE"], errors="coerce")
    return df


def inspections(df):
    """Una fila por inspección (el CSV trae una fila por infracción)."""
    g = df.groupby(["CAMIS", "date", "INSPECTION TYPE"], sort=False)
    ins = g.agg(boro=("BORO", "first"), cuisine=("CUISINE DESCRIPTION", "first"),
                score=("SCORE", "max"), grade=("GRADE", "first"),
                n_viol=("VIOLATION CODE", "count"),
                n_crit=("CRITICAL FLAG", lambda s: (s == "Critical").sum()),
                closed=("ACTION", lambda s: s.str.contains("Closed", na=False).any())).reset_index()
    return ins.sort_values(["CAMIS", "date"]).reset_index(drop=True)


def history_features(ins):
    """Para cada inspección, resumen de las inspecciones ANTERIORES del mismo restaurante."""
    g = ins.groupby("CAMIS", sort=False)
    ins["prev_score"] = g["score"].shift(1)
    ins["prev_crit"] = g["n_crit"].shift(1)
    ins["prev_viol"] = g["n_viol"].shift(1)
    ins["days_since_prev"] = (ins["date"] - g["date"].shift(1)).dt.days
    ins["n_prev"] = g.cumcount()
    ins["prev_closures"] = ins.groupby("CAMIS")["closed"].transform(lambda s: s.cumsum().shift(1).fillna(0))
    ins["prev_mean_score"] = g["score"].transform(lambda s: s.shift(1).expanding().mean())
    ins["prev_max_score"] = g["score"].transform(lambda s: s.shift(1).expanding().max())
    bad = ins["grade"].isin(["B", "C"]).astype(int)
    ins["prev_bad_grades"] = bad.groupby(ins["CAMIS"]).transform(lambda s: s.shift(1).fillna(0).cumsum())
    ins["month"] = ins["date"].dt.month
    return ins


FEATURES = ["prev_score", "prev_crit", "prev_viol", "days_since_prev", "n_prev", "prev_closures",
            "prev_mean_score", "prev_max_score", "prev_bad_grades", "month", "boro_c", "cuisine_c"]


def capture_at(y, score, share):
    """% de suspensos que caen en el `share` de restaurantes con más riesgo."""
    k = int(len(y) * share)
    top = np.argsort(-score)[:k]
    return float(y[top].sum() / y.sum())


def main(path):
    raw = load(path)
    ins = history_features(inspections(raw))

    # ---------- Descriptivo ----------
    graded = ins[ins["grade"].isin(["A", "B", "C"])]
    boro = (graded.groupby("boro")["grade"].value_counts(normalize=True).unstack().fillna(0) * 100).round(1)

    init = ins[(ins["INSPECTION TYPE"] == INITIAL) & ins["score"].notna()].copy()
    init["fail"] = (init["score"] >= FAIL_SCORE).astype(int)
    cuis = init.groupby("cuisine").agg(n=("fail", "size"), tasa=("fail", "mean"))
    cuis = cuis[cuis["n"] >= 800].sort_values("tasa")
    cuis["tasa"] = (cuis["tasa"] * 100).round(1)

    viol = raw.dropna(subset=["VIOLATION CODE"]).copy()
    top_viol = viol.groupby("VIOLATION CODE").agg(n=("CAMIS", "size"), desc=("VIOLATION DESCRIPTION", "first"))
    top_viol = top_viol.sort_values("n", ascending=False).head(8)

    # ---------- Modelo de riesgo (sin fuga) ----------
    top_cuis = init["cuisine"].value_counts().head(40).index
    init["cuisine_c"] = init["cuisine"].where(init["cuisine"].isin(top_cuis), "Other").astype("category").cat.codes
    init["boro_c"] = init["boro"].astype("category").cat.codes
    tr, te = init[init["date"] < SPLIT_DATE], init[init["date"] >= SPLIT_DATE]

    clf = HistGradientBoostingClassifier(max_iter=400, learning_rate=0.05, max_leaf_nodes=31,
                                         categorical_features=[FEATURES.index("boro_c"), FEATURES.index("cuisine_c")],
                                         random_state=42)
    clf.fit(tr[FEATURES], tr["fail"])
    proba = clf.predict_proba(te[FEATURES])[:, 1]
    y = te["fail"].values
    auc = roc_auc_score(y, proba)

    # Regla sencilla de referencia: "suspendió la última vez"
    rule = te["prev_score"].fillna(-1).values
    auc_rule = roc_auc_score(y, rule)

    shares = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
    curve = {"share": [int(s * 100) for s in shares],
             "modelo": [round(capture_at(y, proba, s) * 100, 1) for s in shares],
             "regla": [round(capture_at(y, rule + np.random.default_rng(0).random(len(rule)) * 1e-3, s) * 100, 1)
                       for s in shares]}

    out = {
        "kpis": {
            "registros": int(len(raw)),
            "inspecciones": int(len(ins)),
            "restaurantes": int(ins["CAMIS"].nunique()),
            "iniciales": int(len(init)),
            "tasa_suspenso": round(init["fail"].mean() * 100, 1),
            "auc_fuga": 0.9998,
            "auc_modelo": round(auc, 3),
            "auc_regla": round(auc_rule, 3),
            "captura_30": round(capture_at(y, proba, 0.3) * 100, 1),
            "train_n": int(len(tr)), "test_n": int(len(te)), "split": SPLIT_DATE,
        },
        "boro": {b: {g: float(boro.loc[b, g]) for g in ("A", "B", "C")} for b in boro.index},
        "cocinas": [{"k": k, "n": int(r.n), "tasa": float(r.tasa)} for k, r in cuis.iterrows()],
        "infracciones": [{"code": c, "n": int(r.n), "desc": r.desc[:110]} for c, r in top_viol.iterrows()],
        "captura": curve,
    }
    OUT.parent.mkdir(exist_ok=True)
    OUT.write_text(json.dumps(out, ensure_ascii=False, separators=(",", ":")))
    print(json.dumps(out["kpis"], indent=2), out["boro"], out["cocinas"][:3], out["cocinas"][-3:], curve, sep="\n")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else ROOT / "raw" / "nyc_inspections.csv")