"""Model D — competitor price-response event study / distributed lags (H3, association only).

Outcome: sustained own price cut on day t (S2 listing-days). Treatment: indicator of
>= 1 sustained cut among local competitors (same make-model, year +-1, other dealer,
within 25 km; 10/50 km robustness) on day t-k, k = -7..+14. Linear probability
model with market x week and make-model-year x week fixed effects (DECISIONS D-011),
SEs clustered by dealer (primary) and by region (secondary). Placebo: competitor-cut
series permuted across listings within make-model (geography broken).
"""
from __future__ import annotations

import json
import time

import numpy as np
import pandas as pd
import polars as pl
import pyfixest as pf

from . import config as C

OUT = C.RESULTS / "competition"
LEADS, LAGS = 7, 14


def load(radius: str = "25") -> pd.DataFrame:
    RS = pl.read_parquet(C.PROCESSED / "risk_set.parquet").filter(pl.col("s2_dynamic"))
    L = pl.read_parquet(C.PROCESSED / "listings.parquet").select(["listing_id", "dealer_id", "make", "model", "year", "region", "powertrain", "body_filled", "source_class"])
    RS = RS.join(L, on="listing_id", how="left").sort(["listing_id", "t"])
    RS = RS.with_columns([(pl.col(f"ccut_{radius}") > 0).cast(pl.Int8).alias("cc"), pl.col("first_cut").fill_null(0)])
    shifts = []
    for k in range(-LEADS, LAGS + 1):
        name = f"cc_m{abs(k)}" if k < 0 else f"cc_p{k}"   # m = lead (future competitor cut), p = lag (past)
        shifts.append(pl.col("cc").shift(k).over("listing_id").fill_null(0).alias(name))
    RS = RS.with_columns(shifts)
    RS = RS.with_columns([
        (pl.col("t") // 7).alias("week"),
        pl.concat_str([pl.col("region").fill_null("(missing)"), pl.lit("|"), (pl.col("t") // 7).cast(pl.Utf8)]).alias("mkt_week"),
        pl.concat_str([pl.col("make"), pl.lit("|"), pl.col("model"), pl.lit("|"), pl.col("year").cast(pl.Utf8), pl.lit("|"), (pl.col("t") // 7).cast(pl.Utf8)]).alias("mmy_week"),
        pl.concat_str([pl.col("make"), pl.lit("|"), pl.col("model")]).alias("mm"),
        pl.col("listing_age").clip(upper_bound=60).alias("age_c"),
        (pl.col(f"dens_{radius}") + 1).log().alias("ldens"),
    ])
    df = RS.to_pandas()
    for c in ["mkt_week", "mmy_week", "mm", "dealer_id", "region"]:
        df[c] = pd.factorize(df[c])[0]
    return df


def lag_names():
    return [f"cc_m{k}" for k in range(LEADS, 0, -1)] + [f"cc_p{k}" for k in range(0, LAGS + 1)]


def fit(df: pd.DataFrame, y: str = "cut", fe: str = "mkt_week + mmy_week", cluster: str = "dealer_id", name: str = "main") -> dict:
    t0 = time.time()
    rhs = " + ".join(lag_names()) + " + ldens + C(age_c)"
    m = pf.feols(f"{y} ~ {rhs} | {fe}", data=df, vcov={"CRV1": cluster}, fixef_rm="none")
    t = m.tidy().reset_index(); t = t.rename(columns={t.columns[0]: "term"})
    t = t[t["term"].isin(lag_names())].copy()
    t["k"] = t["term"].map(lambda s: -int(s[4:]) if s.startswith("cc_m") else int(s[4:]))
    t = t.sort_values("k"); t.to_csv(OUT / f"eventstudy_{name}.csv", index=False)
    names = list(m._coefnames)
    def joint(cols):
        idx = [names.index(c) for c in cols]; b = np.asarray(m.coef())[idx]; V = np.asarray(m._vcov)[np.ix_(idx, idx)]
        W = float(b @ np.linalg.solve(V, b)); from scipy.stats import chi2
        return {"stat": W, "df": len(idx), "p": float(chi2.sf(W, len(idx))), "sum": float(b.sum())}
    leads = [f"cc_m{k}" for k in range(1, LEADS + 1)]; lags = [f"cc_p{k}" for k in range(0, LAGS + 1)]
    out = {"n": int(m._N), "base_rate": float(df[y].mean()), "fe": fe, "cluster": cluster, "seconds": round(time.time() - t0, 1),
           "leads_joint": joint(leads), "lags_joint": joint(lags), "lag0_7_sum": joint([f"cc_p{k}" for k in range(0, 8)]),
           "coef": {r.term: {"b": float(r.Estimate), "se": float(r["Std. Error"]), "lo": float(r["2.5%"]), "hi": float(r["97.5%"])} for _, r in t.iterrows()}}
    return out


def placebo(df: pd.DataFrame, seed: int = 36) -> dict:
    """Permute the competitor-cut series across listings within make-model (breaks geography)."""
    rng = np.random.default_rng(seed)
    cols = lag_names()
    first = df.groupby("listing_id")["mm"].first()
    perm = {}
    for mm, ids in first.groupby(first).groups.items():
        ids = np.array(list(ids)); sh = ids.copy(); rng.shuffle(sh); perm.update(dict(zip(ids, sh)))
    # map each listing's series to another listing's series of the same make-model, aligned on listing_age
    src = df[["listing_id", "listing_age"] + cols].copy(); src["listing_id"] = src["listing_id"].map(perm)
    d2 = df.drop(columns=cols).merge(src, on=["listing_id", "listing_age"], how="left")
    d2[cols] = d2[cols].fillna(0)
    return fit(d2, name="placebo_permuted")


def selection_check(df: pd.DataFrame) -> dict:
    """Do listings that ever receive a competitor cut differ at entry? (selection into treatment)"""
    g = df.groupby("listing_id").agg(any_cc=("cc", "max"), ldens=("ldens", "first"), logp=("logp", "first"), pp=("pricepos_25", "first"))
    out = {}
    for v in ["ldens", "logp", "pp"]:
        a, b = g.loc[g.any_cc == 1, v].dropna(), g.loc[g.any_cc == 0, v].dropna()
        out[v] = {"mean_treated": float(a.mean()), "mean_untreated": float(b.mean()), "std_diff": float((a.mean() - b.mean()) / np.sqrt((a.var() + b.var()) / 2))}
    out["share_ever_treated"] = float(g.any_cc.mean())
    return out


def run() -> dict:
    OUT.mkdir(parents=True, exist_ok=True)
    res = {}
    df = load("25")
    print("main", flush=True); res["main_25km"] = fit(df)
    print("cluster region", flush=True); res["main_25km_cluster_region"] = fit(df, cluster="region", name="main_cluster_region")
    print("first cut only", flush=True)
    res["firstcut_25km"] = fit(df[df["in_first_cut_risk"]], y="first_cut", name="firstcut")
    print("no FE / coarse", flush=True); res["coarse_fe"] = fit(df, fe="region + week", name="coarse_fe")
    print("placebo", flush=True); res["placebo_permuted"] = placebo(df)
    res["selection"] = selection_check(df)
    for r in ("10", "50"):
        print("radius", r, flush=True); d = load(r); res[f"main_{r}km"] = fit(d, name=f"main_{r}km")
    json.dump(res, open(OUT / "summary.json", "w"), indent=2, default=float)
    return res


if __name__ == "__main__":
    print(json.dumps(run(), indent=2, default=float))