"""Données de panel et identification causale (chapitre 36) : transformation within, erreurs-types en grappes (CR1),
effets fixes bidirectionnels par indicatrices (via hedonique.mco), doubles différences 2 × 2, étude d'événement,
estimateur de Callaway–Sant'Anna (2021) par cohortes avec contrôles « jamais traités », décomposition de
Goodman-Bacon (2021), contrôle synthétique (Abadie–Diamond–Hainmueller 2010) et discontinuité locale linéaire.

Toutes les fonctions sont pures (numpy/pandas/scipy) et testées dans analysis/python/tests/test_ch36.py.
Conventions : un panel « long » est un DataFrame avec une colonne d'unité, une colonne de période entière (0, 1, …, T−1),
une variable dépendante et, pour les plans échelonnés, une colonne `g` = première période traitée (NaN si jamais traité)."""
from __future__ import annotations
import numpy as np, pandas as pd
from scipy import optimize
from . import hedonique as h

# ----------------------------------------------------------------------------- MCO en grappes et within

def mco_grappes(X: np.ndarray, y: np.ndarray, grappes: np.ndarray) -> dict:
    """MCO avec erreurs-types classiques, HC1 et robustes aux grappes (CR1 : facteur G/(G−1)·(N−1)/(N−K)).
    X contient la constante (ou les indicatrices qui en tiennent lieu)."""
    X = np.asarray(X, float); y = np.asarray(y, float); grappes = np.asarray(grappes)
    base = h.mco(X, y)
    n, k = X.shape; e = base["residus"]
    bread = np.linalg.pinv(X.T @ X)
    ids = np.unique(grappes); G = len(ids)
    meat = np.zeros((k, k))
    for gid in ids:
        m = grappes == gid
        s = X[m].T @ e[m]
        meat += np.outer(s, s)
    c = G / (G - 1) * (n - 1) / (n - k)
    V = c * bread @ meat @ bread
    out = dict(base); out.update({"se_grappes": np.sqrt(np.diag(V)), "V_grappes": V, "G": G, "correction_cr1": c})
    return out

def within(y: np.ndarray, X: np.ndarray, unites: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Transformation within (écarts aux moyennes d'unité) de y et de chaque colonne de X."""
    df = pd.DataFrame(np.column_stack([y, X])); g = df.groupby(np.asarray(unites))
    Z = (df - g.transform("mean")).to_numpy()
    return Z[:, 0], Z[:, 1:]

def within_deux_voies(y: np.ndarray, X: np.ndarray, unites: np.ndarray, temps: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Transformation within bidirectionnelle (panel cylindré) : z − z̄_i − z̄_t + z̄."""
    df = pd.DataFrame(np.column_stack([y, X])); u = np.asarray(unites); t = np.asarray(temps)
    Z = (df - df.groupby(u).transform("mean") - df.groupby(t).transform("mean") + df.mean()).to_numpy()
    return Z[:, 0], Z[:, 1:]

def estimateur_within(y, X, unites, grappes=None) -> dict:
    """Estimateur within (effets fixes d'unité) : MCO sur les données transformées, ddl corrigés du nombre d'unités,
    erreurs-types en grappes par unité (par défaut)."""
    yw, Xw = within(y, X, unites)
    grappes = np.asarray(unites) if grappes is None else np.asarray(grappes)
    r = mco_grappes(Xw, yw, grappes)
    n, k = Xw.shape; N = len(np.unique(unites))
    # correction des degrés de liberté pour les erreurs-types classiques (N constantes éliminées)
    s2 = r["residus"] @ r["residus"] / (n - k - N)
    r["se"] = np.sqrt(np.diag(s2 * np.linalg.pinv(Xw.T @ Xw))); r["sigma2"] = s2
    return r

def indicatrices(codes, omettre_premiere: bool = True, prefixe: str = "") -> pd.DataFrame:
    """Indicatrices (0/1) d'un vecteur de codes, dans l'ordre trié ; la première modalité est omise par défaut."""
    s = pd.Series(np.asarray(codes))
    D = pd.get_dummies(s, prefix=prefixe, prefix_sep="" if prefixe == "" else ":", dtype=float)
    if omettre_premiere: D = D.iloc[:, 1:]
    return D.reset_index(drop=True)

def twfe(df: pd.DataFrame, y: str, regresseurs: list[str], unite: str, temps: str, fe_supp: str | None = None,
         grappes: str | None = None) -> dict:
    """Effets fixes bidirectionnels par indicatrices : y = Σ β_r x_r + α_i + λ_t (+ θ_{s,t} si fe_supp) + ε.
    fe_supp : colonne de strate (ex. RMR) → indicatrices strate × période (période 0 omise dans chaque strate), qui
    remplacent les effets de période communs. Erreurs-types en grappes par `grappes` (défaut : l'unité)."""
    d = df.reset_index(drop=True)
    parts = [pd.DataFrame({"const": np.ones(len(d))}), d[regresseurs].astype(float).reset_index(drop=True),
             indicatrices(d[unite], prefixe="u")]
    if fe_supp is None:
        parts.append(indicatrices(d[temps], prefixe="t"))
    else:
        st = d[fe_supp].astype(str) + "@" + d[temps].astype(int).astype(str)
        D = indicatrices(st, omettre_premiere=False, prefixe="st")
        t0 = d[temps].min()
        D = D[[c for c in D.columns if not c.endswith(f"@{int(t0)}")]]
        parts.append(D.reset_index(drop=True))
    X = pd.concat(parts, axis=1)
    cl = d[unite].values if grappes is None else d[grappes].values
    r = mco_grappes(X.values, d[y].values, cl)
    idx = {nom: i for i, nom in enumerate(X.columns)}
    coef = pd.DataFrame({"coef": [r["beta"][idx[v]] for v in regresseurs], "se": [r["se"][idx[v]] for v in regresseurs],
                         "se_hc1": [r["se_hc1"][idx[v]] for v in regresseurs], "se_grappes": [r["se_grappes"][idx[v]] for v in regresseurs]},
                        index=regresseurs)
    return {"coef": coef, "n": r["n"], "k": r["k"], "G": r["G"], "sigma2": r["sigma2"], "r2": r["r2"], "residus": r["residus"], "colonnes": list(X.columns), "beta": r["beta"]}

# ----------------------------------------------------------------------------- doubles différences 2 × 2

def did_2x2(df: pd.DataFrame, y: str, unite: str, temps: str, traite: str, avant, apres) -> dict:
    """DiD à deux périodes : quatre moyennes de cellule (fenêtres `avant` et `apres` = listes de périodes), estimateur
    (ȳ_tr,ap − ȳ_tr,av) − (ȳ_té,ap − ȳ_té,av), et version agrégée par unité (différence longue Δy_i régressée sur
    l'indicatrice de traitement, erreurs-types HC1 : robustes à l'autocorrélation intra-unité)."""
    d = df[df[temps].isin(list(avant) + list(apres))].copy()
    d["periode"] = np.where(d[temps].isin(list(apres)), 1, 0)
    cell = d.groupby([traite, "periode"])[y].mean().unstack()
    est = (cell.loc[1, 1] - cell.loc[1, 0]) - (cell.loc[0, 1] - cell.loc[0, 0])
    # agrégation par unité : Δy_i = ȳ_i,après − ȳ_i,avant
    u = d.groupby([unite, traite, "periode"])[y].mean().unstack("periode")
    delta = (u[1] - u[0]).reset_index(); delta.columns = [unite, traite, "dy"]
    X = np.column_stack([np.ones(len(delta)), delta[traite].values.astype(float)])
    r = h.mco(X, delta["dy"].values)
    return {"cellules": cell, "estimateur": float(est), "coef_reg": float(r["beta"][1]), "se_hc1": float(r["se_hc1"][1]),
            "se_classique": float(r["se"][1]), "n_traites": int(delta[traite].sum()), "n_temoins": int((delta[traite] == 0).sum()),
            "deltas": delta}

# ----------------------------------------------------------------------------- étude d'événement

def event_study(df: pd.DataFrame, y: str, unite: str, temps: str, g: str, fenetre: tuple[int, int] = (-8, 8), ref: int = -1,
                fe_supp: str | None = None, cohortes=None) -> pd.DataFrame:
    """Étude d'événement par effets fixes bidirectionnels : indicatrices du temps relatif k = t − g (bornes de `fenetre`
    regroupées : « ≤ kmin » et « ≥ kmax »), période de référence `ref` omise, unités jamais traitées (g NaN) à zéro.
    `cohortes` : restreint les unités traitées à ces valeurs de g (les autres cohortes sont exclues du panel)."""
    d = df.copy()
    if cohortes is not None:
        d = d[d[g].isna() | d[g].isin(list(cohortes))].copy()
    kmin, kmax = fenetre
    k = d[temps] - d[g]
    k = k.where(d[g].notna())
    kb = k.clip(lower=kmin, upper=kmax)
    noms = []
    for j in range(kmin, kmax + 1):
        if j == ref: continue
        nom = f"k{j:+d}"; d[nom] = ((kb == j) & d[g].notna()).astype(float); noms.append(nom)
    r = twfe(d, y, noms, unite, temps, fe_supp=fe_supp)
    out = r["coef"].copy(); out["k"] = [int(n[1:]) for n in out.index]
    out = pd.concat([out, pd.DataFrame({"coef": [0.0], "se": [0.0], "se_hc1": [0.0], "se_grappes": [0.0], "k": [ref]}, index=[f"k{ref:+d}"])]).sort_values("k")
    out["ic_bas"] = out["coef"] - 1.96 * out["se_grappes"]; out["ic_haut"] = out["coef"] + 1.96 * out["se_grappes"]
    return out.set_index("k")

# ----------------------------------------------------------------------------- Callaway–Sant'Anna (simplifié)

def callaway_santanna(df: pd.DataFrame, y: str, unite: str, temps: str, g: str, strate: str | None = None,
                      controle: str = "jamais", B: int = 999, graine: int = 20261031) -> dict:
    """ATT(g, t) = E[y_t − y_{g−1} | G = g] − E[y_t − y_{g−1} | C], période de base universelle g − 1, contrôles
    « jamais traités » (controle='jamais') ou « pas encore traités à t » (controle='pas_encore') ; si `strate` est donnée,
    les contrôles sont restreints à la strate de la cohorte (ex. même RMR). Agrégations : dynamique (par k), par cohorte,
    globale (pondérée par la taille des cohortes). Erreurs-types : bootstrap par blocs d'unités (B tirages)."""
    W = df.pivot(index=unite, columns=temps, values=y)
    G = df.groupby(unite)[g].first().reindex(W.index)
    S = df.groupby(unite)[strate].first().reindex(W.index) if strate else None
    periodes = list(W.columns); cohortes = sorted(G.dropna().unique())

    def _att_table(W, G, S):
        rows = []
        for gg in cohortes:
            tr = G.index[G == gg]
            base = int(gg) - 1
            for t in periodes:
                if controle == "jamais":
                    ctl = G.index[G.isna()]
                else:
                    ctl = G.index[G.isna() | (G > max(t, gg))]
                if S is not None:
                    s_g = S.loc[tr].mode().iloc[0]; ctl = [c for c in ctl if S.loc[c] == s_g]
                dy_tr = W.loc[tr, t] - W.loc[tr, base]; dy_c = W.loc[ctl, t] - W.loc[ctl, base]
                att = dy_tr.mean() - dy_c.mean()
                se = np.sqrt(dy_tr.var(ddof=1) / len(dy_tr) + dy_c.var(ddof=1) / len(dy_c)) if len(dy_tr) > 1 and len(dy_c) > 1 else np.nan
                rows.append({"g": int(gg), "t": int(t), "k": int(t) - int(gg), "att": float(att), "se_analytique": float(se), "n_g": len(tr), "n_c": len(ctl)})
        return pd.DataFrame(rows)

    def _agreger(tab):
        post = tab[tab["k"] >= 0]
        par_cohorte = post.groupby("g").apply(lambda d: d["att"].mean(), include_groups=False)
        taille = post.groupby("g")["n_g"].first()
        globale = float((par_cohorte * taille).sum() / taille.sum())
        dyn = tab.groupby("k").apply(lambda d: float((d["att"] * d["n_g"]).sum() / d["n_g"].sum()), include_groups=False)
        return par_cohorte, globale, dyn, taille

    tab = _att_table(W, G, S)
    par_cohorte, globale, dyn, taille = _agreger(tab)
    rng = np.random.default_rng(graine)
    boot_c, boot_g, boot_d = [], [], []
    unites = np.array(W.index)
    for _ in range(B):
        # bootstrap stratifié : chaque cohorte et chaque strate de contrôles (ex. RMR) est rééchantillonnée séparément,
        # de sorte que le plan (tailles des groupes, présence de témoins dans chaque strate) est conservé
        idx = []
        for gg in cohortes:
            u = unites[(G == gg).values]; idx += list(rng.choice(u, size=len(u), replace=True))
        strates_ctl = S.loc[G.isna()].unique() if S is not None else [None]
        for st in strates_ctl:
            m = G.isna().values if st is None else (G.isna() & (S == st)).values
            u = unites[m]; idx += list(rng.choice(u, size=len(u), replace=True))
        Wb = W.loc[idx].reset_index(drop=True); Gb = G.loc[idx].reset_index(drop=True); Sb = S.loc[idx].reset_index(drop=True) if S is not None else None
        Wb.index = range(len(Wb)); Gb.index = range(len(Gb))
        if Sb is not None: Sb.index = range(len(Sb))
        tb = _att_table(Wb, Gb, Sb); pc, gl, dn, _ = _agreger(tb)
        boot_c.append(pc.reindex(par_cohorte.index).values); boot_g.append(gl); boot_d.append(dn.reindex(dyn.index).values)
    boot_c = np.array(boot_c); boot_d = np.array(boot_d)
    cohortes_df = pd.DataFrame({"att": par_cohorte, "se_boot": boot_c.std(axis=0, ddof=1), "n_g": taille})
    dyn_df = pd.DataFrame({"att": dyn, "se_boot": np.nanstd(boot_d, axis=0, ddof=1)})
    dyn_df["ic_bas"] = dyn_df["att"] - 1.96 * dyn_df["se_boot"]; dyn_df["ic_haut"] = dyn_df["att"] + 1.96 * dyn_df["se_boot"]
    return {"att_gt": tab, "par_cohorte": cohortes_df, "globale": globale, "globale_se_boot": float(np.std(boot_g, ddof=1)), "dynamique": dyn_df}

# ----------------------------------------------------------------------------- décomposition de Goodman-Bacon

def bacon(df: pd.DataFrame, y: str, unite: str, temps: str, g: str) -> dict:
    """Décomposition de Goodman-Bacon (2021) de l'estimateur TWFE statique (panel cylindré, sans covariable) en
    comparaisons 2 × 2 : chaque cohorte vs jamais traités ; cohorte précoce (traitée) vs cohorte tardive (témoin, avant
    son traitement) ; cohorte tardive (traitée) vs cohorte précoce (témoin déjà traité). Vérifie Σ s = 1 et
    Σ s·β_2x2 = β_TWFE."""
    W = df.pivot(index=unite, columns=temps, values=y)
    G = df.groupby(unite)[g].first().reindex(W.index)
    T = W.shape[1]; periodes = np.array(W.columns); N = W.shape[0]
    cohortes = sorted(G.dropna().unique())
    n = {gg: (G == gg).sum() / N for gg in cohortes}; n_U = G.isna().sum() / N
    Dbar = {gg: (periodes >= gg).mean() for gg in cohortes}
    # variance de D doublement centré
    d = df.copy(); d["D"] = ((d[temps] >= d[g]) & d[g].notna()).astype(float)
    Dt = d["D"] - d.groupby(unite)["D"].transform("mean") - d.groupby(temps)["D"].transform("mean") + d["D"].mean()
    VD = float((Dt ** 2).mean())

    def m(ids, mask_t):
        return W.loc[ids, periodes[mask_t]].values.mean()
    rows = []
    U = G.index[G.isna()]
    for gg in cohortes:
        tr = G.index[G == gg]
        if len(U) > 0:
            b = (m(tr, periodes >= gg) - m(tr, periodes < gg)) - (m(U, periodes >= gg) - m(U, periodes < gg))
            nkU = n[gg] / (n[gg] + n_U)
            s = (n[gg] + n_U) ** 2 * nkU * (1 - nkU) * Dbar[gg] * (1 - Dbar[gg]) / VD
            rows.append({"traites": int(gg), "temoins": "jamais", "type": "traités vs jamais traités", "beta": float(b), "poids": float(s)})
    for i, gk in enumerate(cohortes):
        for gl in cohortes[i + 1:]:
            k_ids = G.index[G == gk]; l_ids = G.index[G == gl]
            nkl = n[gk] / (n[gk] + n[gl])
            # précoce (k) traité, tardif (l) témoin : fenêtre t < gl
            b_kl = (m(k_ids, (periodes >= gk) & (periodes < gl)) - m(k_ids, periodes < gk)) - (m(l_ids, (periodes >= gk) & (periodes < gl)) - m(l_ids, periodes < gk))
            s_kl = ((n[gk] + n[gl]) * (1 - Dbar[gl])) ** 2 * nkl * (1 - nkl) * ((Dbar[gk] - Dbar[gl]) / (1 - Dbar[gl])) * ((1 - Dbar[gk]) / (1 - Dbar[gl])) / VD
            rows.append({"traites": int(gk), "temoins": int(gl), "type": "précoce vs tardive (pas encore traitée)", "beta": float(b_kl), "poids": float(s_kl)})
            # tardif (l) traité, précoce (k) témoin déjà traité : fenêtre t ≥ gk
            b_lk = (m(l_ids, periodes >= gl) - m(l_ids, (periodes >= gk) & (periodes < gl))) - (m(k_ids, periodes >= gl) - m(k_ids, (periodes >= gk) & (periodes < gl)))
            s_lk = ((n[gk] + n[gl]) * Dbar[gk]) ** 2 * nkl * (1 - nkl) * (Dbar[gl] / Dbar[gk]) * ((Dbar[gk] - Dbar[gl]) / Dbar[gk]) / VD
            rows.append({"traites": int(gl), "temoins": int(gk), "type": "tardive vs précoce (déjà traitée)", "beta": float(b_lk), "poids": float(s_lk)})
    tab = pd.DataFrame(rows)
    tw = twfe(d, y, ["D"], unite, temps)
    return {"table": tab, "somme_poids": float(tab["poids"].sum()), "beta_pondere": float((tab["beta"] * tab["poids"]).sum()),
            "beta_twfe": float(tw["coef"].loc["D", "coef"]), "var_D": VD}

# ----------------------------------------------------------------------------- contrôle synthétique

def controle_synthetique(Y0: np.ndarray, y1: np.ndarray, T0: int) -> dict:
    """Poids w ≥ 0, Σ w = 1, minimisant Σ_{t<T0} (y1_t − Y0_t w)² (Abadie, Diamond et Hainmueller 2010, sans prédicteurs
    auxiliaires). Y0 : (T × J) donneurs ; y1 : (T,) unité traitée. Retourne poids, trajectoire synthétique, écarts,
    RMSPE pré et post."""
    Y0 = np.asarray(Y0, float); y1 = np.asarray(y1, float); J = Y0.shape[1]
    A = Y0[:T0]; b = y1[:T0]
    def f(w): r = A @ w - b; return float(r @ r)
    def grad(w): return 2 * A.T @ (A @ w - b)
    cons = ({"type": "eq", "fun": lambda w: w.sum() - 1.0, "jac": lambda w: np.ones(J)},)
    w0 = np.ones(J) / J
    res = optimize.minimize(f, w0, jac=grad, bounds=[(0, 1)] * J, constraints=cons, method="SLSQP", options={"maxiter": 2000, "ftol": 1e-14})
    w = np.clip(res.x, 0, None); w = w / w.sum()
    synth = Y0 @ w; gap = y1 - synth
    return {"poids": w, "synthetique": synth, "ecarts": gap, "rmspe_pre": float(np.sqrt(np.mean(gap[:T0] ** 2))),
            "rmspe_post": float(np.sqrt(np.mean(gap[T0:] ** 2))), "effet_moyen_post": float(gap[T0:].mean()), "convergence": bool(res.success)}

def placebos_synthetiques(Y0: np.ndarray, T0: int) -> pd.DataFrame:
    """Placebos « dans l'espace » : chaque donneur devient à son tour l'unité traitée, les autres forment le bassin."""
    Y0 = np.asarray(Y0, float); J = Y0.shape[1]; rows = []
    for j in range(J):
        r = controle_synthetique(np.delete(Y0, j, axis=1), Y0[:, j], T0)
        rows.append({"donneur": j, "rmspe_pre": r["rmspe_pre"], "rmspe_post": r["rmspe_post"], "ratio": r["rmspe_post"] / r["rmspe_pre"], "effet_moyen_post": r["effet_moyen_post"]})
    return pd.DataFrame(rows)

# ----------------------------------------------------------------------------- discontinuité (locale linéaire)

def rdd_local_lineaire(x: np.ndarray, y: np.ndarray, h_bande: float, c: float = 0.0) -> dict:
    """Saut à la frontière c par régression locale linéaire à noyau triangulaire de bande h : y = a + τ·1[x ≥ c] + b(x−c)
    + d·1[x ≥ c](x−c), pondérée par K((x−c)/h). Erreur-type HC1 de τ."""
    x = np.asarray(x, float); y = np.asarray(y, float); u = (x - c) / h_bande
    w = np.clip(1 - np.abs(u), 0, None); m = w > 0
    D = (x[m] >= c).astype(float); xc = x[m] - c
    X = np.column_stack([np.ones(m.sum()), D, xc, D * xc]); sw = np.sqrt(w[m])
    r = h.mco(X * sw[:, None], y[m] * sw)
    return {"tau": float(r["beta"][1]), "se_hc1": float(r["se_hc1"][1]), "n_eff": int(m.sum()), "beta": r["beta"], "bande": h_bande}