"""Automobile Market Liquidity Index (AMLI) — Phase 5 (analysis_plan §4, DECISIONS D-013).

Higher AMLI = more liquid. Market-week panel (region x ISO-like week of the panel,
weeks 0..10 complete; week 11 partial is dropped). Components computed with
information available in week w, from the reliable-source listing-day universe (S2L):
  exit      = confirmed exits / listing-days                      (+)
  mdfreq    = sustained price cuts / listing-days                 (+)
  mddepth   = median sustained cut size (%)                       (+)
  disp      = IQR of hedonic residuals among active listings      (-)
  arrival   = new entries (S2) / active stock                     (+)
  stale     = share of active stock with listing age > 30 d       (-)  [5-component variant only: mechanically censored before week 5]
Weights: equal, PCA (first component, signed), supervised ridge on next-week exit (fit weeks 0-5).
"""
from __future__ import annotations

import json
import time

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

from . import config as C

OUT = C.RESULTS / "amli"
WEEKS = 11                      # complete weeks 0..10 (days 0..76)
WEEK_MIN = 2                    # weeks 0-1 excluded: portal start-up (otogo day 6, autousagee day 11), DECISIONS D-015
MIN_LISTING_DAYS = 30
PRIMARY = ["exit", "mdfreq", "mddepth", "disp", "arrival"]
SIGN = {"exit": 1, "mdfreq": 1, "mddepth": 1, "disp": -1, "arrival": 1, "stale": -1}


def listing_days() -> pd.DataFrame:
    RS = pl.read_parquet(C.PROCESSED / "risk_set.parquet").select(["listing_id", "t", "listing_age", "cut", "exit", "s2_dynamic", "left_censored", "logp"])
    L = pl.read_parquet(C.PROCESSED / "listings.parquet").select(["listing_id", "dealer_id", "region", "entry_price", "body_filled"])
    R = pl.read_parquet(C.RESULTS / "hedonic" / "residuals.parquet").select(["listing_id", "resid_A2"])
    E = pl.read_parquet(C.PROCESSED / "price_events.parquet").filter(pl.col("kind") == "cut").with_columns(
        ((pl.col("ts") / 86400).floor() - (C.WINDOW_START_TS // 86400)).cast(pl.Int32).alias("t")).select(["listing_id", "t", "change_pct"])
    d = RS.join(L, on="listing_id", how="left").join(R, on="listing_id", how="left").join(E, on=["listing_id", "t"], how="left")
    d = d.with_columns([(pl.col("t") // 7).alias("week"), (pl.col("listing_age") == 0).cast(pl.Int8).alias("entry"),
                        (pl.col("listing_age") > 30).cast(pl.Int8).alias("stale_flag"), pl.col("region").fill_null("(missing)")])
    # price tier (terciles of entry price over S2L)
    q = d.select(pl.col("entry_price").quantile(0.333).alias("q1"), pl.col("entry_price").quantile(0.667).alias("q2")).row(0)
    d = d.with_columns(pl.when(pl.col("entry_price") < q[0]).then(pl.lit("low")).when(pl.col("entry_price") < q[1]).then(pl.lit("mid")).otherwise(pl.lit("high")).alias("tier"))
    return d.filter((pl.col("week") < WEEKS) & (pl.col("week") >= WEEK_MIN)).to_pandas()


def components(d: pd.DataFrame, by: list[str], w: np.ndarray | None = None) -> pd.DataFrame:
    """Aggregate listing-days to market-week components; w = optional row weights (bootstrap)."""
    d = d.copy()
    d["w"] = 1.0 if w is None else w
    d["w_exit"] = d["w"] * d["exit"]; d["w_cut"] = d["w"] * d["cut"]; d["w_entry"] = d["w"] * d["entry"] * d["s2_dynamic"]
    d["w_stale"] = d["w"] * d["stale_flag"]
    g = d.groupby(by + ["week"])
    out = g.agg(ld=("w", "sum"), exits=("w_exit", "sum"), cuts=("w_cut", "sum"), entries=("w_entry", "sum"), stale=("w_stale", "sum")).reset_index()
    out["exit"] = out["exits"] / out["ld"]; out["mdfreq"] = out["cuts"] / out["ld"]; out["arrival"] = out["entries"] / out["ld"]; out["stale"] = out["stale"] / out["ld"]
    # depth: weighted median of cut sizes; dispersion: weighted IQR of residuals (approximate by unweighted when w is None)
    cutd = d[d["cut"] == 1].dropna(subset=["change_pct"])
    if w is None:
        dep = cutd.groupby(by + ["week"])["change_pct"].median().rename("mddepth").reset_index()
        rr = d.dropna(subset=["resid_A2"]).groupby(by + ["week"])["resid_A2"].agg(lambda s: s.quantile(0.75) - s.quantile(0.25)).rename("disp").reset_index()
    else:
        def wq(s, ww, q):
            o = np.argsort(s.values); cs = np.cumsum(ww.values[o]); return s.values[o][np.searchsorted(cs, q * cs[-1])]
        dep = cutd.groupby(by + ["week"]).apply(lambda x: wq(x["change_pct"], x["w"], 0.5)).rename("mddepth").reset_index()
        rr = d.dropna(subset=["resid_A2"]).groupby(by + ["week"]).apply(lambda x: wq(x["resid_A2"], x["w"], 0.75) - wq(x["resid_A2"], x["w"], 0.25)).rename("disp").reset_index()
    out = out.merge(dep, on=by + ["week"], how="left").merge(rr, on=by + ["week"], how="left")
    out["mddepth"] = -out["mddepth"] * 100     # positive % cut
    out.loc[out["ld"] < MIN_LISTING_DAYS, PRIMARY + ["stale"]] = np.nan
    return out


def build_index(comp: pd.DataFrame, cols: list[str], fit_weeks=range(WEEK_MIN, 6)) -> tuple[pd.DataFrame, dict]:
    Z = comp.copy()
    mu, sd = {}, {}
    for c in cols:
        mu[c] = float(Z[c].mean()); sd[c] = float(Z[c].std()); Z[f"z_{c}"] = SIGN[c] * (Z[c] - mu[c]) / sd[c]
    zc = [f"z_{c}" for c in cols]
    ok = Z[zc].notna().sum(axis=1) >= min(3, len(cols))
    Z["amli_equal"] = Z.loc[ok, zc].mean(axis=1)
    # PCA on complete rows
    comp_rows = Z[zc].dropna()
    X = comp_rows.values - comp_rows.values.mean(0)
    U, S, Vt = np.linalg.svd(X, full_matrices=False)
    v = Vt[0]; v = v * np.sign(v[cols.index("exit")] if "exit" in cols else v[0])
    expl = float(S[0] ** 2 / (S ** 2).sum())
    Z["amli_pca"] = (Z[zc].fillna(Z[zc].mean()) @ v) / np.abs(v).sum()
    Z.loc[~ok, "amli_pca"] = np.nan
    # supervised: ridge of next-week exit on the other components (fit weeks)
    weights = {"equal": {c: 1 / len(cols) for c in cols}, "pca": dict(zip(cols, v.tolist())), "pca_explained": expl}
    if "exit" in cols:
        others = [f"z_{c}" for c in cols if c != "exit"]
        Z = Z.sort_values([*[c for c in Z.columns if c in ("region", "tier", "body_filled", "market")], "week"])
        key = [c for c in ("region", "tier", "body_filled", "market") if c in Z.columns]
        Z["exit_next"] = Z.groupby(key)["exit"].shift(-1) if key else Z["exit"].shift(-1)
        fitrows = Z[Z["week"].isin(list(fit_weeks))].dropna(subset=others + ["exit_next"])
        A = fitrows[others].values; y = fitrows["exit_next"].values; lam = 1.0
        beta = np.linalg.solve(A.T @ A + lam * np.eye(A.shape[1]), A.T @ (y - y.mean()))
        Z["amli_supervised"] = Z[others].fillna(0) @ beta / np.abs(beta).sum()
        weights["supervised"] = dict(zip([c for c in cols if c != "exit"], beta.tolist()))
        weights["supervised_fit_weeks"] = list(fit_weeks)
    return Z, {"mu": mu, "sd": sd, "weights": weights}


def predictive_validation(Z: pd.DataFrame, key: list[str], eval_weeks=range(6, WEEKS)) -> dict:
    """Leave-exit-out index (equal weights on non-exit components) vs next-week exit intensity, with market FE."""
    import pyfixest as pf
    others = [f"z_{c}" for c in PRIMARY if c != "exit"]
    Z = Z.copy(); Z["amli_noexit"] = Z[others].mean(axis=1)
    Z["mkt"] = Z[key].astype(str).agg("|".join, axis=1) if key else "all"
    Z["exit_next"] = Z.groupby("mkt")["exit"].shift(-1)
    d = Z[Z["week"].isin(list(eval_weeks))].dropna(subset=["amli_noexit", "exit_next", "exit"])
    res = {}
    for name, f in {"noexit_fe": "exit_next ~ amli_noexit | mkt", "noexit_fe_lagexit": "exit_next ~ amli_noexit + exit | mkt", "noexit_pooled": "exit_next ~ amli_noexit"}.items():
        m = pf.feols(f, data=d, vcov={"CRV1": "mkt"} if "|" in f or len(d["mkt"].unique()) > 1 else "hetero")
        t = m.tidy().reset_index(); t = t.rename(columns={t.columns[0]: "term"}); r = t[t["term"] == "amli_noexit"].iloc[0]
        res[name] = {"b": float(r["Estimate"]), "se": float(r["Std. Error"]), "lo": float(r["2.5%"]), "hi": float(r["97.5%"]), "n": int(m._N), "r2": float(m._r2)}
    return res


def bootstrap(d: pd.DataFrame, key: list[str], reps: int = 200, seed: int = 36) -> pd.DataFrame:
    """Block bootstrap over dealers: index values per market-week (equal weights, 5 components)."""
    rng = np.random.default_rng(seed)
    dealers = d["dealer_id"].unique(); idx = pd.factorize(d["dealer_id"])[0]
    rows = []
    for r in range(reps):
        draw = rng.integers(0, len(dealers), len(dealers)); mult = np.bincount(draw, minlength=len(dealers)).astype(float)
        w = mult[idx]
        comp = components(d[w > 0], key, w[w > 0])
        Z, _ = build_index(comp, PRIMARY)
        Z["rep"] = r; rows.append(Z[key + ["week", "amli_equal", "amli_pca", "exit"]])
        if r % 20 == 0:
            print("boot", r, flush=True)
    B = pd.concat(rows)
    q = B.groupby(key + ["week"]).agg(amli_equal_lo=("amli_equal", lambda s: s.quantile(0.025)), amli_equal_hi=("amli_equal", lambda s: s.quantile(0.975)),
                                      amli_pca_lo=("amli_pca", lambda s: s.quantile(0.025)), amli_pca_hi=("amli_pca", lambda s: s.quantile(0.975))).reset_index()
    return q


def run(reps: int = 200) -> dict:
    OUT.mkdir(parents=True, exist_ok=True); t0 = time.time()
    d = listing_days()
    out = {"listing_days": int(len(d)), "weeks": WEEKS}
    # --- region x week (primary) --------------------------------------------
    comp = components(d, ["region"])
    Z, meta = build_index(comp, PRIMARY)
    Z.to_csv(OUT / "index_region_week.csv", index=False); out["region"] = meta
    out["region_validation"] = predictive_validation(Z, ["region"])
    # 5-component variant with staleness (weeks >= 6)
    Z5, meta5 = build_index(comp[comp["week"] >= 6], PRIMARY + ["stale"]); Z5.to_csv(OUT / "index_region_week_5comp.csv", index=False); out["region_5comp"] = meta5
    # leave-one-component-out
    loo = {}
    for c in PRIMARY:
        cols = [x for x in PRIMARY if x != c]; Zc, mc = build_index(comp, cols)
        loo[c] = {"corr_with_full_equal": float(Zc["amli_equal"].corr(Z["amli_equal"])), "pca_explained": mc["weights"]["pca_explained"]}
    out["leave_one_out"] = loo
    out["corr_equal_pca"] = float(Z["amli_equal"].corr(Z["amli_pca"])); out["corr_equal_supervised"] = float(Z["amli_equal"].corr(Z["amli_supervised"]))
    # --- province x week, daily province -------------------------------------
    d["market"] = "Quebec"; compP = components(d, ["market"]); ZP, metaP = build_index(compP, PRIMARY); ZP.to_csv(OUT / "index_province_week.csv", index=False); out["province"] = metaP
    # --- segments -----------------------------------------------------------
    for seg in ["tier", "body_filled"]:
        dd = d.dropna(subset=[seg]); cs = components(dd, [seg]); Zs, ms = build_index(cs, PRIMARY); Zs.to_csv(OUT / f"index_{seg}_week.csv", index=False)
        out[f"segment_{seg}"] = {"pca_explained": ms["weights"]["pca_explained"], "validation": predictive_validation(Zs, [seg])}
    # --- component correlations ------------------------------------------------
    out["component_corr"] = comp[PRIMARY].corr().round(3).to_dict()
    # --- bootstrap --------------------------------------------------------------
    if reps:
        q = bootstrap(d, ["region"], reps=reps); q.to_csv(OUT / "bootstrap_region_week.csv", index=False)
    out["seconds"] = round(time.time() - t0, 1)
    json.dump(out, open(OUT / "summary.json", "w"), indent=2, default=float)
    return out


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