#!/usr/bin/env python3
"""Figures du chapitre 36 — données de panel et identification causale (cas 4, DONNÉES SYNTHÉTIQUES, graine 20261031).
Sorties (assets/figures/ch36/) :
  ch36-fig-trajectoires.pdf    ln RNE moyen (cumul depuis 2016T1) des cohortes traitées et des témoins de la même RMR
  ch36-fig-event-study.pdf     étude d'événement TWFE (RMR × trimestre) et agrégation dynamique de Callaway–Sant'Anna, avec la vérité
  ch36-fig-bacon.pdf           décomposition de Goodman-Bacon : cas 4 et laboratoire à effets croissants
  ch36-fig-synthetique.pdf     contrôle synthétique de la cohorte A : trajectoires, écarts et placebos, poids des donneurs
  ch36-fig-discontinuite.pdf   discontinuité spatiale à une frontière de zone (simulation transversale)
  ch36-fig-sensibilite.pdf     étude d'événement lorsque les tendances parallèles sont violées (+1 pt/an pour les traités)
Usage (racine du dépôt) : PYTHONPATH=analysis/python python3 analysis/python/figures/ch36_figures.py"""
import sys, pathlib
ROOT = pathlib.Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT / "analysis/python")); sys.path.insert(0, str(ROOT / "analysis/python/cases"))
import numpy as np, pandas as pd
import matplotlib.pyplot as plt
from traite.style import setup, palette, savefig, LARGEUR
import ch36_panel as c36

TRIMS = pd.period_range("2016Q1", "2026Q2", freq="Q")
def _v(x, d=3): return f"{x:.{d}f}".replace(".", ",").replace("-", "−")
def _etiq(ax, pas=8):
    ax.set_xticks(range(0, len(TRIMS), pas)); ax.set_xticklabels([str(TRIMS[i]).replace("Q", "T") for i in range(0, len(TRIMS), pas)], fontsize=7.5)

def fig_trajectoires(res):
    se = res["panel"]
    fig, axes = plt.subplots(1, 2, figsize=(LARGEUR, 3.0), sharey=True)
    for ax, (rmr, coh, T0, titre) in zip(axes, [(c36.RMR_A, "A", c36.T_A, "Montréal : cohorte A (2022T3)"), (c36.RMR_B, "B", c36.T_B, "Gatineau : cohorte B (2024T3)")]):
        for i, (lab, m) in enumerate([("Immeubles traités", (se["cohorte"] == coh)), ("Témoins jamais traités, même RMR", (se["cohorte"] == "jamais") & (se["rmr"] == rmr))]):
            W = se[m].pivot(index="id_immeuble", columns="t", values="y")
            traj = (W.sub(W[0], axis=0)).mean(axis=0)
            ax.plot(traj.index, traj.values, color=palette.cat[i], lw=1.7, label=f"{lab} (n = {W.shape[0]})")
            # contrefactuel déclaré : trajectoire traitée moins l'effet vrai moyen
            if i == 0:
                d = se[m].groupby("t")["delta_vrai"].mean()
                ax.plot(traj.index, traj.values - d.values, color=palette.cat[0], lw=1.0, ls=":", label="Traités sans traitement, Y(0) (connu ici)")
        ax.axvline(T0, color=palette.ink2, lw=0.9, ls="--"); ax.text(T0 + 0.4, 0.0, "mise en service", fontsize=7, color=palette.ink2, rotation=90, va="bottom")
        ax.set_title(titre, fontsize=8.5); _etiq(ax); ax.grid(True, lw=0.4)
    axes[0].set_ylabel("ln RNE moyen, écart à 2016T1"); axes[0].legend(fontsize=6.8, loc="upper left")
    savefig(fig, "ch36/ch36-fig-trajectoires.pdf")

def fig_event_study(res):
    es = res["es"]; cs = res["cs"]["dynamique"].loc[-8:15]
    fig, axes = plt.subplots(1, 2, figsize=(LARGEUR, 3.1), sharey=True)
    ax = axes[0]
    ax.fill_between(es.index, es["ic_bas"], es["ic_haut"], color=palette.cat[0], alpha=0.18, lw=0)
    ax.plot(es.index, es["coef"], "o-", color=palette.cat[0], ms=3.8, lw=1.4, label="Coefficient estimé (IC à 95 %, grappes)")
    ax.step(es.index, es["vrai"], where="mid", color=palette.ink, lw=1.1, ls="--", label="Effet déclaré δ(k)")
    ax.set_title("(a) Étude d'événement TWFE (RMR × trimestre)", fontsize=8); ax.set_xlabel("Trimestres depuis la mise en service, k")
    ax.set_ylabel("Effet sur ln RNE (points de log)")
    ax = axes[1]
    ax.fill_between(cs.index, cs["ic_bas"], cs["ic_haut"], color=palette.cat[1], alpha=0.18, lw=0)
    ax.plot(cs.index, cs["att"], "s-", color=palette.cat[1], ms=3.4, lw=1.4, label="ATT(k) de Callaway–Sant'Anna (IC bootstrap)")
    kk = np.arange(-8, 16); ax.step(kk, c36.effet_vrai(kk), where="mid", color=palette.ink, lw=1.1, ls="--", label="Effet déclaré δ(k)")
    ax.set_title("(b) Callaway–Sant'Anna (témoins de la même RMR)", fontsize=8); ax.set_xlabel("Trimestres depuis la mise en service, k")
    for ax in axes:
        ax.axhline(0, color=palette.ink2, lw=0.8); ax.axvline(-0.5, color=palette.ink2, lw=0.8, ls=":"); ax.legend(fontsize=6.8, loc="upper left"); ax.grid(True, lw=0.4)
    savefig(fig, "ch36/ch36-fig-event-study.pdf")

def fig_bacon(res):
    fig, axes = plt.subplots(1, 2, figsize=(LARGEUR, 3.0))
    marqueurs = {"traités vs jamais traités": "o", "précoce vs tardive (pas encore traitée)": "^", "tardive vs précoce (déjà traitée)": "D"}
    for ax, (b, vrai, titre) in zip(axes, [(res["bacon"], res["verite"]["att_global"], "Cas 4 : cohortes A et B, 35 témoins"), (res["labo"]["bacon"], res["labo"]["verite"], "Laboratoire : effets croissants, 10 témoins")]):
        t = b["table"]
        for j, (typ, mk) in enumerate(marqueurs.items()):
            s = t[t["type"] == typ]
            ax.scatter(s["poids"], s["beta"], marker=mk, s=55, color=palette.cat[j], edgecolor=palette.ink, lw=0.5, label=typ, zorder=5)
        ax.axhline(b["beta_twfe"], color=palette.ink, lw=1.1, label="TWFE = Σ poids × (2 × 2)" if ax is axes[0] else None)
        ax.axhline(vrai, color=palette.ink2, lw=1.0, ls="--", label="Effet moyen vrai (ATT)" if ax is axes[0] else None)
        ax.axhline(0, color=palette.grid, lw=0.8)
        ax.text(0.99, b["beta_twfe"], f"TWFE {_v(b['beta_twfe'])}", fontsize=7, color=palette.ink, ha="right", va="bottom", transform=ax.get_yaxis_transform())
        ax.text(0.99, vrai, f"vrai {_v(vrai)}", fontsize=7, color=palette.ink2, ha="right", va="top" if vrai < b["beta_twfe"] else "bottom", transform=ax.get_yaxis_transform())
        ax.set_xlabel("Poids de la comparaison 2 × 2"); ax.set_title(titre, fontsize=8.5); ax.grid(True, lw=0.4)
    axes[0].set_ylabel("Estimation 2 × 2 (points de log)")
    h_, l_ = axes[0].get_legend_handles_labels()
    fig.legend(h_, l_, loc="lower center", ncol=3, fontsize=6.8, bbox_to_anchor=(0.5, -0.1), frameon=False)
    savefig(fig, "ch36/ch36-fig-bacon.pdf")

def fig_synthetique(res):
    s = res["sc"]; T0 = s["T0"]; t = np.arange(len(s["traite"]))
    fig, axes = plt.subplots(1, 3, figsize=(LARGEUR, 2.9), gridspec_kw={"width_ratios": [1.3, 1.3, 1.0]})
    ax = axes[0]
    ax.plot(t, s["traite"], color=palette.cat[0], lw=1.7, label="Cohorte A (moyenne des 8 immeubles)")
    ax.plot(t, s["synthetique"], color=palette.cat[1], lw=1.5, ls="--", label="Contrôle synthétique")
    ax.plot(t, s["traite"] - np.concatenate([np.zeros(T0), c36.effet_vrai(np.arange(len(t) - T0))]), color=palette.cat[0], lw=0.9, ls=":", label="Y(0) déclaré")
    ax.axvline(T0 - 0.5, color=palette.ink2, lw=0.9, ls="--"); ax.set_title("(a) Trajectoires (ln RNE centré)", fontsize=8.5); _etiq(ax, 10); ax.legend(fontsize=6.3, loc="lower right"); ax.grid(True, lw=0.4)
    ax = axes[1]
    W = res["panel"].pivot(index="id_immeuble", columns="t", values="y"); coh = res["panel"].groupby("id_immeuble")["cohorte"].first()
    donneurs = W.loc[coh[coh == "jamais"].index]; Y0 = (donneurs.values - donneurs.values[:, :T0].mean(axis=1, keepdims=True)).T
    from traite import panel as pn
    for j in range(Y0.shape[1]):
        r = pn.controle_synthetique(np.delete(Y0, j, axis=1), Y0[:, j], T0)
        ax.plot(t, r["ecarts"], color=palette.grid, lw=0.7, alpha=0.9, label="Placebos (35 donneurs)" if j == 0 else None)
    ax.plot(t, s["ecarts"], color=palette.cat[0], lw=1.8, label="Cohorte A")
    ax.plot(t, np.concatenate([np.zeros(T0), c36.effet_vrai(np.arange(len(t) - T0))]), color=palette.ink, lw=1.0, ls="--", label="Effet déclaré")
    ax.axvline(T0 - 0.5, color=palette.ink2, lw=0.9, ls="--"); ax.axhline(0, color=palette.ink2, lw=0.7)
    ax.set_title("(b) Écarts et placebos", fontsize=8.5); _etiq(ax, 10); ax.legend(fontsize=6.3, loc="upper left"); ax.grid(True, lw=0.4); ax.set_ylim(-0.12, 0.12)
    ax = axes[2]
    w = s["poids"].sort_values(ascending=True); w = w[w > 0.005]
    ax.barh(range(len(w)), w.values, color=palette.cat[1], height=0.7)
    ax.set_yticks(range(len(w))); ax.set_yticklabels(w.index, fontsize=7); ax.set_xlabel("Poids"); ax.set_title(f"(c) Poids > 0,5 % ({len(w)} donneurs)", fontsize=8.5); ax.grid(True, axis="x", lw=0.4)
    savefig(fig, "ch36/ch36-fig-synthetique.pdf")

def fig_discontinuite(res):
    r = res["rdd"]; x, y = r["x"], r["y"]
    fig, ax = plt.subplots(figsize=(LARGEUR, 3.0))
    ax.scatter(x, y, s=6, color=palette.grid, lw=0, label="Immeubles simulés (n = 400)")
    bords = np.arange(-2, 2.01, 0.2); centres = (bords[:-1] + bords[1:]) / 2
    moy = [y[(x >= a) & (x < b)].mean() for a, b in zip(bords[:-1], bords[1:])]
    ax.plot(centres, moy, "o", color=palette.cat[0], ms=4.5, label="Moyenne par tranche de 200 m")
    e = r["estimations"][1.0]; b = e["beta"]
    xg = np.linspace(-1, 0, 50); ax.plot(xg, b[0] + b[2] * xg, color=palette.cat[1], lw=1.8, label="Ajustement local linéaire (h = 1 km)")
    xd = np.linspace(0, 1, 50); ax.plot(xd, b[0] + b[1] + (b[2] + b[3]) * xd, color=palette.cat[1], lw=1.8)
    ax.axvline(0, color=palette.ink2, lw=0.9, ls="--"); ax.text(0.03, y.min() + 0.01, "frontière de zone", fontsize=7.5, color=palette.ink2)
    ax.annotate("", xy=(0, b[0] + b[1]), xytext=(0, b[0]), arrowprops=dict(arrowstyle="<->", color=palette.ink, lw=1.0))
    ax.annotate(f"saut estimé τ̂ = {_v(e['tau'])} (vrai 0,050)", xy=(0, b[0] + b[1] / 2), xytext=(0.35, 12.72), fontsize=7.5, color=palette.ink, arrowprops=dict(arrowstyle="-", color=palette.ink2, lw=0.6))
    ax.set_xlabel("Distance signée à la frontière (km ; positif = zone de densité supérieure)"); ax.set_ylabel("ln de la valeur simulée")
    ax.legend(fontsize=7, loc="upper left"); ax.grid(True, lw=0.4)
    savefig(fig, "ch36/ch36-fig-discontinuite.pdf")

def fig_sensibilite(res):
    es0 = res["es"]; esv = res["sens"]["event_study"]
    fig, ax = plt.subplots(figsize=(LARGEUR, 3.0))
    ax.fill_between(esv.index, esv["ic_bas"], esv["ic_haut"], color=palette.cat[1], alpha=0.15, lw=0)
    ax.plot(es0.index, es0["coef"], "o-", color=palette.cat[0], ms=3.6, lw=1.3, label="Tendances parallèles respectées (générateur)")
    ax.plot(esv.index, esv["coef"], "s-", color=palette.cat[1], ms=3.6, lw=1.3, label="Traités avec +1 point de croissance par an (IC à 95 %)")
    ax.step(es0.index, es0["vrai"], where="mid", color=palette.ink, lw=1.1, ls="--", label="Effet déclaré δ(k)")
    pre = esv.loc[-7:-2]; p = np.polyfit(pre.index.values.astype(float), pre["coef"].values, 1)
    xx = np.linspace(-8, -1, 20); ax.plot(xx, np.polyval(p, xx), color=palette.cat[1], lw=0.9, ls=":")
    ax.text(-7.9, -0.1, f"pente des coefficients pré-traitement : {_v(p[0], 4)} par trimestre", fontsize=7, color=palette.ink)
    ax.axhline(0, color=palette.ink2, lw=0.8); ax.axvline(-0.5, color=palette.ink2, lw=0.8, ls=":")
    ax.set_xlabel("Trimestres depuis la mise en service, k"); ax.set_ylabel("Effet sur ln RNE (points de log)"); ax.legend(fontsize=7, loc="upper left"); ax.grid(True, lw=0.4)
    savefig(fig, "ch36/ch36-fig-sensibilite.pdf")

def main():
    setup(); res = c36.main(verbose=False)
    for f in (fig_trajectoires, fig_event_study, fig_bacon, fig_synthetique, fig_discontinuite, fig_sensibilite):
        f(res); print("ok", f.__name__)

if __name__ == "__main__":
    main()