"""Siniestralidad vial en Madrid 2024 — análisis exploratorio y modelo de cribado de alcohol.

Genera los agregados que usa la web (data/madrid.json) a partir del parte de
accidentes publicado por el Ayuntamiento de Madrid (datos abiertos, año 2024).
Las definiciones (limpieza, franjas, accidente grave, modelo) son las mismas que
las del notebook del proyecto.

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

import numpy as np
import pandas as pd
from pyproj import Transformer
from scipy.stats import chi2_contingency
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score, precision_score, recall_score
from sklearn.model_selection import train_test_split

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "data" / "madrid.json"

WEEKDAYS = ["Lun", "Mar", "Mié", "Jue", "Vie", "Sáb", "Dom"]
FRANJAS = ["Madrugada", "Mañana", "Tarde", "Noche"]  # 0-5, 6-11, 12-19, 20-23


def load(path):
    df = pd.read_csv(path, sep=";", encoding="utf-8-sig", dtype=str)
    df["fecha"] = pd.to_datetime(df["fecha"], format="%d/%m/%Y", errors="coerce")
    for c in ("coordenada_x_utm", "coordenada_y_utm"):
        df[c] = pd.to_numeric(df[c].str.replace(",", "."), errors="coerce")

    # Limpieza: el tipo de accidente y las coordenadas son del siniestro y no se pueden inventar
    df = df.dropna(subset=["fecha", "num_expediente", "tipo_accidente", "coordenada_x_utm", "coordenada_y_utm"])
    df["estado_meteorológico"] = (df["estado_meteorológico"].fillna("Desconocido")
                                  .replace({"Se desconoce": "Desconocido", "LLuvia intensa": "Lluvia intensa"}))
    df["tipo_vehiculo"] = df["tipo_vehiculo"].fillna("Desconocido")
    # Lesividad en blanco = sin asistencia sanitaria (código 14, según la documentación de la Policía Municipal)
    df["cod_lesividad"] = pd.to_numeric(df["cod_lesividad"], errors="coerce").fillna(14).astype(int)

    df["dia_semana"] = df["fecha"].dt.day_of_week
    df["tipo_dia"] = np.where(df["dia_semana"] >= 5, "Fin de Semana", "Lunes a Viernes")
    df["hora_num"] = pd.to_datetime(df["hora"], format="%H:%M:%S", errors="coerce").dt.hour
    df["franja"] = pd.cut(df["hora_num"], bins=[-1, 5, 11, 19, 23], labels=FRANJAS)
    # Los nulos de la prueba de alcohol se dejan como NaN: solo se excluyen en los análisis de alcohol
    df["positiva_alcohol"] = df["positiva_alcohol"].map({"S": 1, "N": 0})
    # Grave = ingreso hospitalario de más de 24 h (3) o fallecido en 24 h (4)
    df["grave"] = df["cod_lesividad"].isin([3, 4]).astype(int)
    top_vehiculos = df["tipo_vehiculo"].value_counts().head(8).index
    df["grupo_vehiculo"] = df["tipo_vehiculo"].where(df["tipo_vehiculo"].isin(top_vehiculos), "Otros")
    return df


def per_accident(df):
    """Una fila por accidente (el CSV trae una fila por persona implicada)."""
    return df.groupby("num_expediente").agg(
        fecha=("fecha", "first"), dow=("dia_semana", "first"), tipo_dia=("tipo_dia", "first"),
        hora=("hora_num", "first"), franja=("franja", "first"), distrito=("distrito", "first"),
        tipo=("tipo_accidente", "first"), clima=("estado_meteorológico", "first"),
        x=("coordenada_x_utm", "first"), y=("coordenada_y_utm", "first"),
        personas=("num_expediente", "size"), grave=("grave", "max"),
    )


def rate_table(frame, col, target, min_n, title=True):
    t = frame.groupby(col, observed=True)[target].agg(n="size", k="sum")
    t = t[t["n"] >= min_n]
    t["tasa"] = (t["k"] / t["n"] * 100).round(2)
    t = t.sort_values("tasa", ascending=False)
    return [{"k": (k.title() if title and isinstance(k, str) else k), "n": int(r.n), "graves": int(r.k), "tasa": float(r.tasa)}
            for k, r in t.iterrows()]


def chi2_p(frame, col, target):
    return float(chi2_contingency(pd.crosstab(frame[col], frame[target]))[1])


def alcohol_model(df):
    """Regresión logística (nivel persona) que predice el positivo en alcohol, sin fuga de datos."""
    data = df.dropna(subset=["positiva_alcohol"])
    features = ["franja", "tipo_accidente", "rango_edad", "sexo", "grupo_vehiculo"]
    X = pd.get_dummies(data[features].astype(str), dtype=int)  # one-hot: no usa el objetivo
    y = data["positiva_alcohol"].astype(int)
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
    model = LogisticRegression(random_state=42, class_weight="balanced", max_iter=1000).fit(X_tr, y_tr)
    pred = model.predict(X_te)
    tn, fp, fn, tp = confusion_matrix(y_te, pred).ravel()
    coefs = pd.Series(model.coef_[0], index=X.columns).sort_values(ascending=False)
    return {
        "accuracy": round(accuracy_score(y_te, pred) * 100, 1),
        "precision": round(precision_score(y_te, pred) * 100, 1),
        "recall": round(recall_score(y_te, pred) * 100, 1),
        "f1": round(f1_score(y_te, pred), 2),
        "baseline_accuracy": round((1 - y_te.mean()) * 100, 1),
        "cm": {"tn": int(tn), "fp": int(fp), "fn": int(fn), "tp": int(tp)},
        "test_n": int(len(y_te)), "test_pos": int(y_te.sum()),
        "top": [{"k": k, "coef": round(float(v), 2)} for k, v in coefs.head(6).items()],
    }


def main(path):
    df = load(path)
    acc = per_accident(df)
    alc = df.dropna(subset=["positiva_alcohol"])

    # Mapa: accidentes graves en lat/lon (UTM ETRS89 huso 30N -> WGS84)
    geo = acc[(acc["grave"] == 1) & (acc["x"] > 400000) & (acc["y"] > 4400000)]
    tr = Transformer.from_crs("EPSG:25830", "EPSG:4326", always_xy=True)
    lon, lat = tr.transform(geo["x"].values, geo["y"].values)
    points = [[round(a, 5), round(o, 5)] for a, o in zip(lat, lon)]

    # Volumen: matriz día x hora (nº de accidentes)
    heat = (acc.groupby(["dow", "hora"]).size().unstack(fill_value=0)
            .reindex(index=range(7), columns=range(24), fill_value=0))
    peak = heat.stack().idxmax()

    # Gravedad por franja (y por franja y tipo de día), con test chi-cuadrado
    franja = rate_table(acc, "franja", "grave", 0, title=False)
    franja = sorted(franja, key=lambda r: FRANJAS.index(r["k"]))
    franja_dia = {d: sorted(rate_table(acc[acc["tipo_dia"] == d], "franja", "grave", 0, title=False),
                            key=lambda r: FRANJAS.index(r["k"])) for d in ("Lunes a Viernes", "Fin de Semana")}

    # Lluvia frente a no lluvia (excluyendo "Desconocido", granizo y nieve)
    clima_ok = acc[acc["clima"].isin(["Despejado", "Nublado", "Lluvia débil", "Lluvia intensa"])].copy()
    clima_ok["lluvia"] = clima_ok["clima"].str.startswith("Lluvia").astype(int)

    # Alcohol (nivel persona con prueba registrada)
    alc_franja_dia = (alc.groupby(["tipo_dia", "franja"], observed=True)["positiva_alcohol"].mean() * 100).round(2)
    madrugada = alc[alc["franja"] == "Madrugada"]["positiva_alcohol"].mean() * 100

    # Gravedad por tipo de persona (nivel persona)
    persona = (df.groupby("tipo_persona")["grave"].mean() * 100).round(2)

    out = {
        "kpis": {
            "accidentes": int(len(acc)),
            "personas": int(len(df)),
            "graves": int(acc["grave"].sum()),
            "fallecidos": int((df["cod_lesividad"] == 4).sum()),
            "tasa_grave": round(acc["grave"].mean() * 100, 2),
            "alcohol_global": round(alc["positiva_alcohol"].mean() * 100, 2),
            "alcohol_madrugada": round(madrugada, 1),
            "p_franja": chi2_p(acc, "franja", "grave"),
            "p_lluvia": round(chi2_p(clima_ok, "lluvia", "grave"), 2),
            "peaton_grave": float(persona.get("Peatón", np.nan)),
            "conductor_grave": float(persona.get("Conductor", np.nan)),
            # % de personas con lesión grave entre quienes dan positivo y quienes no
            "grave_si_alcohol": round(alc.loc[alc["positiva_alcohol"] == 1, "grave"].mean() * 100, 2),
            "grave_no_alcohol": round(alc.loc[alc["positiva_alcohol"] == 0, "grave"].mean() * 100, 2),
            "pico": {"dow": WEEKDAYS[peak[0]], "hora": int(peak[1]), "n": int(heat.values.max())},
        },
        "mensual": acc.groupby(acc["fecha"].dt.month).size().reindex(range(1, 13), fill_value=0).tolist(),
        "heat": {"dias": WEEKDAYS, "valores": heat.values.tolist()},
        "franja": franja,
        "franja_dia": franja_dia,
        "alcohol_hora": [round(v, 2) for v in
                         (alc.groupby("hora_num")["positiva_alcohol"].mean() * 100).reindex(range(24), fill_value=0)],
        "alcohol_franja_dia": {d: [float(alc_franja_dia.get((d, f), 0)) for f in FRANJAS]
                               for d in ("Lunes a Viernes", "Fin de Semana")},
        "alcohol_tipo": rate_table(alc, "tipo_accidente", "positiva_alcohol", 150),
        "tipos": rate_table(acc, "tipo", "grave", 150),
        "distritos": rate_table(acc, "distrito", "grave", 200),
        "clima": rate_table(acc, "clima", "grave", 50),
        "modelo": alcohol_model(df),
        "puntos_graves": points,
    }
    OUT.parent.mkdir(exist_ok=True)
    OUT.write_text(json.dumps(out, ensure_ascii=False, separators=(",", ":")))
    print(json.dumps(out["kpis"], ensure_ascii=False, indent=2))
    for k in ("franja", "franja_dia", "alcohol_franja_dia", "alcohol_tipo", "clima", "modelo"):
        print(k, json.dumps(out[k], ensure_ascii=False)[:600])
    print("tipos", out["tipos"][:3], "puntos", len(points))


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