"""Chapitre 36 — tests des fonctions de traite.panel (récupération d'effets connus sur données simulées) et des chiffres cités
dans le texte, les tableaux et les solutions. Tolérances : coefficients ±0,0005 (ou relatives déclarées), erreurs-types ±0,0005,
identités (Bacon, within = indicatrices) à 1e-8."""
import sys, pathlib, math
import numpy as np, pandas as pd, pytest
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "cases"))
from traite import panel as pn
from traite import hedonique as h
import ch36_panel as c36

# ----------------------------------------------------------------------------- fonctions génériques (données simulées)

def _panel(N=30, T=12, seed=3, delta=0.05, dyn=0.0):
    rng = np.random.default_rng(seed); alpha = rng.normal(10, 0.5, N); lam = np.cumsum(rng.normal(0.01, 0.01, T))
    g = np.array([np.nan] * 10 + [4] * 10 + [8] * 10, float); rows = []
    for i in range(N):
        for t in range(T):
            k = t - g[i] if not np.isnan(g[i]) else np.nan
            d = (delta + dyn * k) if (not np.isnan(k) and k >= 0) else 0.0
            rows.append({"u": i, "t": t, "g": g[i], "vrai": d, "y": alpha[i] + lam[t] + d + rng.normal(0, 0.01)})
    df = pd.DataFrame(rows); df["D"] = ((df["t"] >= df["g"]) & df["g"].notna()).astype(float); return df

def test_within_egal_indicatrices_et_mco_grappes_reduit_a_hc1():
    df = _panel()
    r_tw = pn.twfe(df, "y", ["D"], "u", "t")
    yw, Xw = pn.within_deux_voies(df["y"].values, df[["D"]].values, df["u"].values, df["t"].values)
    bw = np.linalg.lstsq(Xw, yw, rcond=None)[0][0]
    assert abs(bw - r_tw["coef"].loc["D", "coef"]) < 1e-8
    X = np.column_stack([np.ones(len(df)), df["D"].values]); r = pn.mco_grappes(X, df["y"].values, np.arange(len(df)))
    assert np.allclose(r["se_grappes"], r["se_hc1"], rtol=1e-6)        # une grappe par observation = HC1

def test_twfe_recupere_effet_constant():
    df = _panel(delta=0.05); r = pn.twfe(df, "y", ["D"], "u", "t")
    assert abs(r["coef"].loc["D", "coef"] - 0.05) < 0.005

def test_did_2x2_identite_cellules_et_regression():
    df = _panel(); df["tr"] = (df["g"] == 4).astype(int); d = df[df["g"] != 8]
    r = pn.did_2x2(d, "y", "u", "t", "tr", [1, 2, 3], [4, 5, 6])
    assert abs(r["estimateur"] - r["coef_reg"]) < 1e-10 and abs(r["estimateur"] - 0.05) < 0.01

def test_event_study_retrouve_le_profil():
    df = _panel(delta=0.02, dyn=0.01); es = pn.event_study(df, "y", "u", "t", "g", fenetre=(-3, 3))
    for k in range(0, 3):
        assert abs(es.loc[k, "coef"] - (0.02 + 0.01 * k)) < 0.01
    assert all(abs(es.loc[k, "coef"]) < 0.01 for k in (-3, -2))

def test_bacon_identite_et_cs_corrige_le_twfe():
    df = _panel(N=30, T=16, delta=0.0, dyn=0.02)         # effets croissants : TWFE biaisé
    b = pn.bacon(df, "y", "u", "t", "g")
    assert abs(b["somme_poids"] - 1) < 1e-8 and abs(b["beta_pondere"] - b["beta_twfe"]) < 1e-8
    interdite = b["table"].loc[b["table"]["type"].str.startswith("tardive"), "beta"].iloc[0]
    vrai = df.loc[df["D"] == 1, "vrai"].mean()
    cs = pn.callaway_santanna(df, "y", "u", "t", "g", B=50)
    assert abs(cs["globale"] - vrai) < 0.01 and abs(b["beta_twfe"] - vrai) > 0.02 and interdite < 0

def test_controle_synthetique_retrouve_combinaison_exacte():
    rng = np.random.default_rng(1); T, J, T0 = 30, 6, 20
    Y0 = np.cumsum(rng.normal(0, 0.02, (T, J)), axis=0); w = np.array([0.5, 0.3, 0.2, 0, 0, 0])
    y1 = Y0 @ w + np.concatenate([np.zeros(T0), 0.05 * np.ones(T - T0)])
    r = pn.controle_synthetique(Y0, y1, T0)
    assert np.allclose(r["poids"], w, atol=0.02) and abs(r["effet_moyen_post"] - 0.05) < 0.005 and r["rmspe_pre"] < 1e-3

def test_rdd_retrouve_le_saut():
    rng = np.random.default_rng(2); x = rng.uniform(-1, 1, 4000); y = 1 + 0.2 * x + 0.1 * (x >= 0) + rng.normal(0, 0.05, 4000)
    r = pn.rdd_local_lineaire(x, y, 0.5)
    assert abs(r["tau"] - 0.1) < 3 * r["se_hc1"] and abs(r["tau"] - 0.1) < 0.02

# ----------------------------------------------------------------------------- chiffres du chapitre

@pytest.fixture(scope="module")
def res():
    return c36.main(verbose=False)

def test_plan_et_verite(res):
    v = res["verite"]; tr = res["traites"]
    assert tr["A"] == ["P04", "P06", "P11", "P14", "P20", "P21", "P28", "P42"] and tr["B"] == ["P10", "P24", "P29", "P33", "P37"]
    assert len(res["panel"]) == 2016 and res["panel"]["id_immeuble"].nunique() == 48
    assert v["att_global"] == pytest.approx(0.0454, abs=0.0001) and v["att_A"] == pytest.approx(0.0484, abs=0.0001)
    assert v["att_B"] == pytest.approx(0.0356, abs=0.0001) and v["att_fenetre_2x2"] == pytest.approx(0.04625, abs=1e-6)
    assert c36.effet_vrai(3) == 0.04 and c36.effet_vrai(15) == pytest.approx(0.07) and c36.effet_vrai(-1) == 0
    assert 100 * (math.exp(0.04) - 1) == pytest.approx(4.08, abs=0.005)

def test_did_2x2(res):
    r = res["did_rmr"]; c = r["cellules"]
    assert c.loc[1, 0] == pytest.approx(10.1379, abs=0.0001) and c.loc[1, 1] == pytest.approx(10.2450, abs=0.0001)
    assert c.loc[0, 0] == pytest.approx(10.3696, abs=0.0001) and c.loc[0, 1] == pytest.approx(10.4234, abs=0.0001)
    assert r["estimateur"] == pytest.approx(0.0534, abs=0.0001) and r["se_hc1"] == pytest.approx(0.0201, abs=0.0001)
    assert r["n_traites"] == 8 and r["n_temoins"] == 9
    t = res["did_tous"]; assert t["estimateur"] == pytest.approx(0.0468, abs=0.0001) and t["se_hc1"] == pytest.approx(0.0177, abs=0.0001) and t["n_temoins"] == 35
    p = res["placebo"]; assert p["estimateur"] == pytest.approx(0.0155, abs=0.0001) and p["se_hc1"] == pytest.approx(0.0101, abs=0.0001)

def test_twfe_et_erreurs_types(res):
    tw = res["twfe"]; c = tw["commun"]["coef"].loc["D"]; r = tw["regional"]["coef"].loc["D"]
    assert c["coef"] == pytest.approx(0.0905, abs=0.0001) and tw["within_deux_voies"] == pytest.approx(c["coef"], abs=1e-8)
    assert c["se"] == pytest.approx(0.0047, abs=0.0001) and c["se_hc1"] == pytest.approx(0.0053, abs=0.0001) and c["se_grappes"] == pytest.approx(0.0198, abs=0.0001)
    assert tw["commun"]["k"] == 90 and tw["regional"]["k"] == 295
    assert r["coef"] == pytest.approx(0.0702, abs=0.0001) and r["se"] == pytest.approx(0.0054, abs=0.0001)
    assert r["se_hc1"] == pytest.approx(0.0064, abs=0.0001) and r["se_grappes"] == pytest.approx(0.0264, abs=0.0001)
    assert r["se_grappes"] / r["se"] == pytest.approx(4.9, abs=0.05) and c["se_grappes"] / c["se"] == pytest.approx(4.2, abs=0.05)

def test_event_study(res):
    es = res["es"]
    attendu = {-8: (-0.0175, 0.0204), -7: (-0.0169, 0.0173), -6: (-0.0182, 0.0148), -5: (-0.0121, 0.0131), -4: (-0.0214, 0.0140),
               -3: (-0.0183, 0.0116), -2: (-0.0073, 0.0074), 0: (0.0088, 0.0060), 1: (0.0235, 0.0084), 2: (0.0312, 0.0108),
               3: (0.0410, 0.0138), 4: (0.0354, 0.0141), 5: (0.0439, 0.0149), 6: (0.0476, 0.0155), 7: (0.0561, 0.0158), 8: (0.0817, 0.0270)}
    for k, (b, s) in attendu.items():
        assert es.loc[k, "coef"] == pytest.approx(b, abs=0.0001) and es.loc[k, "se_grappes"] == pytest.approx(s, abs=0.0001)
    assert es.loc[-1, "coef"] == 0 and es.loc[8, "vrai"] == pytest.approx(0.0612, abs=0.0001) and es.loc[4, "vrai"] == pytest.approx(0.0425)
    pre = es.loc[-7:-2]; pente = np.polyfit(pre.index.values.astype(float), pre["coef"].values, 1)[0]
    assert pente == pytest.approx(0.0011, abs=0.0001)

def test_bacon_cas4(res):
    b = res["bacon"]; t = b["table"].set_index("type")
    assert b["somme_poids"] == pytest.approx(1, abs=1e-10) and b["beta_pondere"] == pytest.approx(b["beta_twfe"], abs=1e-10)
    assert b["beta_twfe"] == pytest.approx(0.0905, abs=0.0001)
    tj = b["table"][b["table"]["temoins"] == "jamais"].set_index("traites")
    assert tj.loc[26, "poids"] == pytest.approx(0.6658, abs=0.0001) and tj.loc[26, "beta"] == pytest.approx(0.0844, abs=0.0001)
    assert tj.loc[34, "poids"] == pytest.approx(0.2721, abs=0.0001) and tj.loc[34, "beta"] == pytest.approx(0.1224, abs=0.0001)
    assert t.loc["précoce vs tardive (pas encore traitée)", "poids"] == pytest.approx(0.0476, abs=0.0001) and t.loc["précoce vs tardive (pas encore traitée)", "beta"] == pytest.approx(0.0080, abs=0.0001)
    assert t.loc["tardive vs précoce (déjà traitée)", "poids"] == pytest.approx(0.0146, abs=0.0001) and t.loc["tardive vs précoce (déjà traitée)", "beta"] == pytest.approx(0.0403, abs=0.0001)
    # vérité de la comparaison A vs B avant 2024T3 : moyenne de δ(k) pour k = 0..7
    assert c36.effet_vrai(np.arange(0, 8)).mean() == pytest.approx(0.0356, abs=0.0001)

def test_callaway_santanna(res):
    c = res["cs"]
    assert c["globale"] == pytest.approx(0.0481, abs=0.0001) and c["globale_se_boot"] == pytest.approx(0.0116, abs=0.0002)
    pc = c["par_cohorte"]
    assert pc.loc[26, "att"] == pytest.approx(0.0528, abs=0.0001) and pc.loc[26, "se_boot"] == pytest.approx(0.0154, abs=0.0002)
    assert pc.loc[34, "att"] == pytest.approx(0.0406, abs=0.0001) and pc.loc[34, "se_boot"] == pytest.approx(0.0183, abs=0.0002)
    d = c["dynamique"]
    assert d.loc[3, "att"] == pytest.approx(0.0417, abs=0.0001) and d.loc[15, "att"] == pytest.approx(0.0740, abs=0.0001)
    c2 = res["cs_tous"]
    assert c2["globale"] == pytest.approx(0.0514, abs=0.0001) and c2["par_cohorte"].loc[26, "att"] == pytest.approx(0.0509, abs=0.0001)
    assert c2["par_cohorte"].loc[34, "att"] == pytest.approx(0.0521, abs=0.0001)

def test_controle_synthetique_cas4(res):
    s = res["sc"]
    assert s["effet_moyen_post"] == pytest.approx(0.0363, abs=0.0001) and s["vrai_post"] == pytest.approx(0.0484, abs=0.0001)
    assert s["rmspe_pre"] == pytest.approx(0.0031, abs=0.0001) and s["rmspe_post"] == pytest.approx(0.0397, abs=0.0001)
    assert s["ratio_traite"] == pytest.approx(12.9, abs=0.1) and s["rang_ratio"] == 1 and len(s["placebos"]) == 35
    w = s["poids"]; assert abs(w.sum() - 1) < 1e-8 and (w >= -1e-12).all() and int((w > 0.005).sum()) == 11
    for k, v in {"P02": 0.142, "P09": 0.136, "P16": 0.123, "P25": 0.120, "P19": 0.114, "P43": 0.110, "P45": 0.105, "P03": 0.071}.items():
        assert w[k] == pytest.approx(v, abs=0.001)

def test_discontinuite(res):
    r = res["rdd"]; e = r["estimations"]
    assert r["saut_vrai"] == 0.05 and r["naif"] == pytest.approx(0.1136, abs=0.0001)
    assert e[0.5]["tau"] == pytest.approx(0.034, abs=0.0005) and e[0.5]["se_hc1"] == pytest.approx(0.0343, abs=0.0001) and e[0.5]["n_eff"] == 93
    assert e[1.0]["tau"] == pytest.approx(0.0316, abs=0.0001) and e[1.0]["se_hc1"] == pytest.approx(0.0247, abs=0.0001) and e[1.0]["n_eff"] == 202
    assert e[2.0]["tau"] == pytest.approx(0.058, abs=0.0005) and e[2.0]["se_hc1"] == pytest.approx(0.0183, abs=0.0001) and e[2.0]["n_eff"] == 400
    for hb in e: assert abs(e[hb]["tau"] - 0.05) < 1.0 * e[hb]["se_hc1"]

def test_laboratoire(res):
    L = res["labo"]
    assert L["twfe"] == pytest.approx(0.0233, abs=0.0001) and L["se_grappes"] == pytest.approx(0.0061, abs=0.0001)
    assert L["verite"] == pytest.approx(0.0692, abs=0.0001) and L["verite_k"] == pytest.approx(0.0825, abs=0.0001) and L["verite_l"] == pytest.approx(0.0425, abs=0.0001)
    assert L["cs"]["globale"] == pytest.approx(0.0601, abs=0.0001) and L["cs"]["globale_se_boot"] == pytest.approx(0.0073, abs=0.0002)
    pc = L["cs"]["par_cohorte"]; assert pc.loc[10, "att"] == pytest.approx(0.0893, abs=0.0001) and pc.loc[26, "att"] == pytest.approx(0.0310, abs=0.0001)
    t = L["bacon"]["table"].set_index("type")
    assert t.loc["tardive vs précoce (déjà traitée)", "beta"] == pytest.approx(-0.0356, abs=0.0001) and t.loc["tardive vs précoce (déjà traitée)", "poids"] == pytest.approx(0.3265, abs=0.0001)
    assert t.loc["précoce vs tardive (pas encore traitée)", "beta"] == pytest.approx(0.0404, abs=0.0001)
    assert abs(L["bacon"]["beta_pondere"] - L["twfe"]) < 1e-8 and abs(L["twfe"] - L["verite"]) / L["se_grappes"] > 7

def test_sensibilite_et_temoins(res):
    sv = res["sens"]
    assert sv["did_2x2"]["estimateur"] == pytest.approx(0.0734, abs=0.0001) and sv["twfe_regional"]["coef"].loc["D", "coef"] == pytest.approx(0.1227, abs=0.0001)
    assert sv["pente_pre"] == pytest.approx(0.0036, abs=0.0001) and sv["twfe_tendance"]["coef"].loc["D", "coef"] == pytest.approx(0.0609, abs=0.0001)
    assert sv["did_detendance"] == pytest.approx(0.0426, abs=0.0001) and sv["ecart_ans"] == 2.0
    es = sv["event_study"]; assert es.loc[-8, "coef"] == pytest.approx(-0.0609, abs=0.0001) and es.loc[7, "coef"] == pytest.approx(0.0761, abs=0.0001)
    # biais du 2 × 2 = tendance × écart entre fenêtres (0,01 × 2 ans)
    assert sv["did_2x2"]["estimateur"] - res["did_rmr"]["estimateur"] == pytest.approx(0.02, abs=1e-6)
    t = res["temoins"].set_index("temoins")
    assert t.loc["jamais traités, toutes RMR", "estimateur"] == pytest.approx(0.0468, abs=0.0001) and t.loc["jamais traités, toutes RMR", "n_temoins"] == 35
    assert t.loc["jamais traités, Montréal", "estimateur"] == pytest.approx(0.0534, abs=0.0001)
    assert t.loc["jamais traités, hors Montréal", "estimateur"] == pytest.approx(0.0445, abs=0.0001) and t.loc["jamais traités, hors Montréal", "n_temoins"] == 26
    assert t.loc["pas encore traités (cohorte B)", "estimateur"] == pytest.approx(0.0308, abs=0.0001) and t.loc["pas encore traités (cohorte B)", "n_temoins"] == 5

# ----------------------------------------------------------------------------- solutions

def test_solutions_exercices(res):
    ex = res["ex"]
    assert ex["ex1"]["did"] == pytest.approx(0.05) and ex["ex1"]["pct"] == pytest.approx(5.13, abs=0.005)
    assert np.allclose(ex["ex2"]["yw"], [-0.15, 0.15, -0.05, 0.05, -0.05, 0.05]) and np.allclose(ex["ex2"]["xw"], [-0.5, 0.5, -0.5, 0.5, 0, 0])
    assert ex["ex2"]["b_within"] == pytest.approx(0.20) and ex["ex2"]["b_pooled"] == pytest.approx(0.425)
    assert ex["ex3"][(42, 0.5)] == pytest.approx(4.64, abs=0.005) and ex["ex3"][(42, 0.9)] == pytest.approx(6.16, abs=0.005) and ex["ex3"][(8, 0.5)] == pytest.approx(2.12, abs=0.005)
    assert (4.9 ** 2 - 1) / 41 == pytest.approx(0.56, abs=0.005)
    e5 = ex["ex5"]; assert e5["Dk"] == 0.75 and e5["Dl"] == 0.25 and e5["VD"] == pytest.approx(0.0625) and e5["s_kl"] == pytest.approx(0.5) and e5["s_lk"] == pytest.approx(0.5)
    # (c) espérances : β_kl = 0,055 ; β_lk = 0,03 − (0,13 − 0,055) = −0,045 ; TWFE = 0,005 ; ATT global 0,0675
    att_k = 0.01 * np.arange(1, 16).mean(); att_l = 0.01 * np.arange(1, 6).mean()
    assert att_k == pytest.approx(0.08) and att_l == pytest.approx(0.03)
    b_kl = 0.01 * np.arange(1, 11).mean(); b_lk = att_l - (0.01 * np.arange(11, 16).mean() - b_kl)
    assert b_kl == pytest.approx(0.055) and b_lk == pytest.approx(-0.045) and 0.5 * b_kl + 0.5 * b_lk == pytest.approx(0.005)
    assert (15 * att_k + 5 * att_l) / 20 == pytest.approx(0.0675)
    e6 = ex["ex6"]; assert e6["att"][2] == pytest.approx(0.045) and e6["att"][3] == pytest.approx(0.0817, abs=0.0001) and e6["att"][0] == pytest.approx(-0.015) and e6["moyenne_post"] == pytest.approx(0.0633, abs=0.0001)

def test_solution_exercice_7(res):
    t = res["ex7"].set_index("variante")
    assert t.loc["35 donneurs, six RMR", "effet_post"] == pytest.approx(0.0363, abs=0.0001) and t.loc["35 donneurs, six RMR", "n_poids_sup_1pct"] == 11
    assert t.loc["9 donneurs, Montréal", "rmspe_pre"] == pytest.approx(0.0067, abs=0.0001) and t.loc["9 donneurs, Montréal", "effet_post"] == pytest.approx(0.0452, abs=0.0001)
    assert t.loc["9 donneurs, Montréal", "n_poids_sup_1pct"] == 7 and t.loc["9 donneurs, Montréal", "trois_premiers"].startswith("P09 0.32, P27 0.20")
    assert t.loc["26 donneurs, hors Montréal", "rmspe_pre"] == pytest.approx(0.0035, abs=0.0001) and t.loc["26 donneurs, hors Montréal", "effet_post"] == pytest.approx(0.0375, abs=0.0001)
    assert t.loc["placebo 2020T3, 35 donneurs", "rmspe_pre"] == pytest.approx(0.0027, abs=0.0001) and t.loc["placebo 2020T3, 35 donneurs", "effet_post"] == pytest.approx(-0.0055, abs=0.0001)