"""LaTeX table fragments (booktabs tabular only) for model results: tab04 hedonic, tab05 hazards,
tab06 competition/spatial, tab07 AMLI. Tables 1-3 come from descriptives.py, 8 from prediction.py,
9 from robustness.py. Every number is read from results/."""
from __future__ import annotations

import json

import numpy as np
import pandas as pd

from . import config as C

T = C.TABLES
R = C.RESULTS


def _f(x, d=3):
    return "" if x is None or (isinstance(x, float) and np.isnan(x)) else f"{x:.{d}f}"


def _stars(p):
    return "$^{***}$" if p < 0.01 else "$^{**}$" if p < 0.05 else "$^{*}$" if p < 0.10 else ""


def _n(x):
    return f"{int(x):,}".replace(",", "\\,")


def _write(name, lines):
    T.mkdir(parents=True, exist_ok=True)
    (T / name).write_text("\n".join(lines) + "\n")


def tab04_hedonic():
    S = json.load(open(R / "hedonic" / "summary.json")); cv = json.load(open(R / "hedonic" / "cv_grouped_dealer.json"))["summary"]
    specs = [("A1_linear_mm", "A1 linear"), ("A2_spline_mm", "A2 splines"), ("A2t_tensor_mm", "A2t tensor"), ("A3_spline_x_powertrain", "A3 splines $\\times$ powertrain"), ("A4_spline_mmt", "A4 splines, trim FE"), ("A5_spline_dealerFE", "A5 + dealer FE")]
    coefs = {k: pd.read_csv(R / "hedonic" / f"coefs_{k}.csv").set_index("term") for k, _ in specs}
    rows = [("age", "Vehicle age (years)"), ("km", "Mileage (10\\,000 km)"), ("C(powertrain)[T.BEV]", "Battery electric"), ("C(powertrain)[T.HEV]", "Hybrid"), ("C(powertrain)[T.PHEV]", "Plug-in hybrid"), ("C(powertrain)[T.Diesel]", "Diesel"),
            ("C(transmission)[T.Manual]", "Manual transmission"), ("C(drivetrain)[T.FWD]", "Front-wheel drive"), ("C(drivetrain)[T.RWD]", "Rear-wheel drive"), ("C(source_class)[T.full_portal]", "Full-inventory portal"), ("C(source_class)[T.banner]", "Dealer group")]
    L = ["\\begin{tabular}{l" + "c" * len(specs) + "}", "\\toprule", "& " + " & ".join(f"({i+1})" for i in range(len(specs))) + " \\\\", "& " + " & ".join(lab for _, lab in specs) + " \\\\", "\\midrule"]
    for term, lab in rows:
        cells, ses = [], []
        for k, _ in specs:
            c = coefs[k]
            if k == "A5_spline_dealerFE" and term.startswith("C(source_class)"):
                cells.append("absorbed"); ses.append(""); continue
            if term in c.index:
                r = c.loc[term]; cells.append(_f(r["Estimate"]) + _stars(r["Pr(>|t|)"])); ses.append("(" + _f(r["Std. Error"]) + ")")
            else:
                cells.append("--"); ses.append("")
        L.append(lab + " & " + " & ".join(cells) + " \\\\"); L.append(" & " + " & ".join(ses) + " \\\\")
    L += ["\\midrule", "Nonlinear age/km terms & no & yes & yes & yes & yes & yes \\\\", "Tensor age $\\times$ km & no & no & yes & no & no & no \\\\",
          "Quality FE & make-model & make-model & make-model & make-model & make-model-trim & make-model-trim \\\\", "Entry-week FE & yes & yes & yes & yes & yes & yes \\\\", "Region dummies & yes & yes & yes & yes & yes & absorbed \\\\", "Dealer FE & no & no & no & no & no & yes \\\\"]
    L.append("Observations & " + " & ".join(_n(S["fits"][k]["n"]) for k, _ in specs) + " \\\\")
    L.append("$R^2$ & " + " & ".join(_f(S["fits"][k]["r2"]) for k, _ in specs) + " \\\\")
    L.append("Within $R^2$ & " + " & ".join(_f(S["fits"][k]["r2_within"]) for k, _ in specs) + " \\\\")
    L.append("In-sample RMSE (log) & " + " & ".join(_f(S["fits"][k]["rmse"]) for k, _ in specs) + " \\\\")
    L.append("Grouped-dealer CV RMSE & " + " & ".join(_f(cv[k]["rmse_mean"]) if k in cv else "--" for k, _ in specs) + " \\\\")
    L += ["\\bottomrule", "\\end{tabular}"]
    _write("tab04_hedonic.tex", L)
    # tests fragment
    t = S["tests"]
    L2 = ["\\begin{tabular}{lccc}", "\\toprule", "Test & $\\chi^2$ & df & $p$ \\\\", "\\midrule"]
    for k, lab in [("H1_nonlinear_age", "H1: nonlinear age terms = 0 (A2)"), ("H1_nonlinear_km", "H1: nonlinear mileage terms = 0 (A2)"), ("H1_nonlinear_joint", "H1: all nonlinear terms = 0 (A2)"), ("H1_tensor_interaction", "H1: age $\\times$ mileage tensor = 0 (A2t)"), ("H1_age_x_powertrain", "H1: age profile $\\times$ powertrain = 0 (A3)"), ("H4_region_joint", "H4: region effects = 0 (A4)")]:
        v = t[k]; L2.append(f"{lab} & {v['stat']:.1f} & {v['df']} & {v['p']:.3g} \\\\")
    L2 += ["\\bottomrule", "\\end{tabular}"]
    _write("tab04b_hedonic_tests.tex", L2)


def tab05_hazards():
    H = json.load(open(R / "hazards" / "summary.json"))
    B = pd.read_csv(R / "hazards" / "cloglog_B_first_markdown.csv").set_index("term"); Cx = pd.read_csv(R / "hazards" / "cloglog_C_exit.csv").set_index("term")
    cb = pd.read_csv(R / "hazards" / "cox_first_markdown.csv").set_index("term"); ce = pd.read_csv(R / "hazards" / "cox_exit.csv").set_index("term")
    cs = pd.read_csv(R / "hazards" / "cutsize_ols.csv").set_index("term")
    rows = [("ldens", "ldens_entry", "ldens", "log(1 + competitors within 25 km)"), ("resid_A2", "resid_A2", "resid_A2", "Entry hedonic residual (relative price)"), ("pricepos_entry", "pricepos_entry", None, "Price position vs competitors at entry"),
            ("pricepos_25", None, "pricepos_25", "Current price position vs competitors"), ("prior_cuts", None, None, "Prior sustained cuts (count)"), ("lprice", "lprice", "lprice", "log entry price"), ("lkm", "lkm", "lkm", "log mileage"), ("age_years_entry", "age_years_entry", "age_years_entry", "Vehicle age at entry (years)"), ("own_stock", None, None, "Own-dealer same-model stock")]
    L = ["\\begin{tabular}{lccccc}", "\\toprule", "& \\multicolumn{2}{c}{First markdown} & \\multicolumn{2}{c}{Listing exit} & Cut size \\\\", "\\cmidrule(lr){2-3}\\cmidrule(lr){4-5}\\cmidrule(lr){6-6}",
         "& cloglog & Cox & cloglog & Cox & log(\\%) \\\\", "\\midrule"]
    def cell(df, term, est="coef", se="se", p="p"):
        if term is None or term not in df.index:
            return "--", ""
        r = df.loc[term]; return _f(r[est]) + _stars(r[p]), "(" + _f(r[se]) + ")"
    for tB, tC, tS, lab in rows:
        c1 = cell(B, tB); c2 = cell(cb, tC, "coef", "se(coef)", "p"); c3 = cell(Cx, tB); c4 = cell(ce, tC, "coef", "se(coef)", "p"); c5 = cell(cs, tS, "Estimate", "Std. Error", "Pr(>|t|)")
        L.append(lab + " & " + " & ".join(x[0] for x in (c1, c2, c3, c4, c5)) + " \\\\"); L.append(" & " + " & ".join(x[1] for x in (c1, c2, c3, c4, c5)) + " \\\\")
    L += ["\\midrule", f"Listing-age bins (joint $\\chi^2$, df 8) & {H['B_first_markdown_cloglog']['joint_age_bins']['stat']:.1f} & -- & {H['C_exit_cloglog']['joint_age_bins']['stat']:.1f} & -- & -- \\\\",
          f"Listing-days / listings & {_n(H['B_first_markdown_cloglog']['n_rows'])} & {_n(H['B_cut_cox']['n'])} & {_n(H['C_exit_cloglog']['n_rows'])} & {_n(H['C_exit_cox']['n'])} & {_n(H['B_cut_size']['n'])} \\\\",
          f"Events & {_n(H['B_first_markdown_cloglog']['n_events'])} & {_n(H['B_cut_cox']['events'])} & {_n(H['C_exit_cloglog']['n_events'])} & {_n(H['C_exit_cox']['events'])} & -- \\\\",
          f"Concordance / within $R^2$ & -- & {_f(H['B_cut_cox']['concordance'])} & -- & {_f(H['C_exit_cox']['concordance'])} & {_f(H['B_cut_size']['r2_within'])} \\\\",
          f"PH test: covariates with $p<0.01$ & -- & {H['B_cut_cox']['ph_violations_p<0.01']} & -- & {H['C_exit_cox']['ph_violations_p<0.01']} & -- \\\\",
          "\\bottomrule", "\\end{tabular}"]
    _write("tab05_hazards.tex", L)


def tab06_competition():
    S = json.load(open(R / "competition" / "summary.json")); S2 = json.load(open(R / "competition" / "summary_dow.json")) if (R / "competition" / "summary_dow.json").exists() else {}
    specs = [("main_25km", "Main, 25 km"), ("main_25km_cluster_region", "Cluster: region"), ("main_10km", "10 km"), ("main_50km", "50 km"), ("coarse_fe", "Coarse FE"), ("placebo_permuted", "Placebo")]
    specs2 = [("dealer_dow_fe", "+ dealer $\\times$ weekday FE"), ("dealer_dow_fe_firstcut", "First cut, + dealer $\\times$ weekday")]
    allspecs = [(k, lab, S[k]) for k, lab in specs] + [(k, lab, S2[k]) for k, lab in specs2 if k in S2]
    L = ["\\begin{tabular}{l" + "c" * len(allspecs) + "}", "\\toprule", "& " + " & ".join(lab for _, lab, _ in allspecs) + " \\\\", "\\midrule"]
    for k, lab in [("cc_m7", "$k=-7$"), ("cc_m3", "$k=-3$"), ("cc_m1", "$k=-1$"), ("cc_p0", "$k=0$"), ("cc_p1", "$k=+1$"), ("cc_p2", "$k=+2$"), ("cc_p3", "$k=+3$"), ("cc_p7", "$k=+7$"), ("cc_p14", "$k=+14$")]:
        L.append(lab + " & " + " & ".join(f"{100*v['coef'][k]['b']:.2f}" for _, _, v in allspecs) + " \\\\")
        L.append(" & " + " & ".join(f"({100*v['coef'][k]['se']:.2f})" for _, _, v in allspecs) + " \\\\")
    L += ["\\midrule", "Sum of leads $k=-7..-1$ & " + " & ".join(f"{100*v['leads_joint']['sum']:.2f}" for _, _, v in allspecs) + " \\\\",
          "\\quad joint $p$ & " + " & ".join(f"{v['leads_joint']['p']:.2g}" for _, _, v in allspecs) + " \\\\",
          "Sum of lags $k=0..7$ & " + " & ".join(f"{100*v['lag0_7_sum']['sum']:.2f}" for _, _, v in allspecs) + " \\\\",
          "\\quad joint $p$ & " + " & ".join(f"{v['lag0_7_sum']['p']:.2g}" for _, _, v in allspecs) + " \\\\",
          "Daily base rate (\\%) & " + " & ".join(f"{100*v['base_rate']:.2f}" for _, _, v in allspecs) + " \\\\",
          "Listing-days & " + " & ".join(_n(v["n"]) for _, _, v in allspecs) + " \\\\", "\\bottomrule", "\\end{tabular}"]
    _write("tab06_competition.tex", L)
    # spatial: region premia + dispersion
    prem = pd.read_csv(R / "hedonic" / "region_premia_A4.csv"); p2 = pd.read_csv(R / "hedonic" / "region_premia_A2.csv").set_index("region"); disp = pd.read_csv(R / "hedonic" / "dispersion_by_region.csv").set_index("region")
    prem = prem[prem["region"] != "(missing)"].sort_values("Estimate", ascending=False)
    L = ["\\begin{tabular}{lcccc}", "\\toprule", "Region & A4 premium (\\%) & 95\\% CI & A2 premium (\\%) & Residual IQR \\\\", "\\midrule"]
    for _, r in prem.iterrows():
        L.append(f"{r['region']} & {100*r['Estimate']:.1f} & [{100*r['2.5%']:.1f}, {100*r['97.5%']:.1f}] & {100*p2.loc[r['region'], 'Estimate']:.1f} & {disp.loc[r['region'], 'iqr']:.3f} \\\\")
    L.append(f"Montréal (reference) & 0 & -- & 0 & {disp.loc['Montréal', 'iqr']:.3f} \\\\")
    L += ["\\bottomrule", "\\end{tabular}"]
    _write("tab06b_region_premia.tex", L)


def tab07_amli():
    A = json.load(open(R / "amli" / "summary.json"))
    w = A["region"]["weights"]; comps = ["exit", "mdfreq", "mddepth", "disp", "arrival"]; lab = {"exit": "Exit intensity", "mdfreq": "Markdown frequency", "mddepth": "Markdown depth", "disp": "Residual dispersion (sign reversed)", "arrival": "Arrival rate"}
    L = ["\\begin{tabular}{lcccc}", "\\toprule", "Component & Equal & PCA loading & Supervised & Leave-one-out corr. \\\\", "\\midrule"]
    for c in comps:
        sup = w["supervised"].get(c); L.append(f"{lab[c]} & {w['equal'][c]:.2f} & {w['pca'][c]:+.2f} & {'--' if sup is None else f'{sup:+.4f}'} & {A['leave_one_out'][c]['corr_with_full_equal']:.2f} \\\\")
    L += ["\\midrule", f"PCA first-component share & \\multicolumn{{4}}{{c}}{{{A['region']['weights']['pca_explained']:.2f}}} \\\\", f"Corr(equal, PCA) / Corr(equal, supervised) & \\multicolumn{{4}}{{c}}{{{A['corr_equal_pca']:.2f} / {A['corr_equal_supervised']:.2f}}} \\\\"]
    v = A["region_validation"]
    for k, l in [("noexit_fe", "Leave-exit-out AMLI $\\to$ next-week exit (region FE)"), ("noexit_fe_lagexit", "\\quad + lagged exit"), ("noexit_pooled", "\\quad pooled OLS")]:
        L.append(f"{l} & \\multicolumn{{4}}{{c}}{{{100*v[k]['b']:.2f} pp [{100*v[k]['lo']:.2f}, {100*v[k]['hi']:.2f}], $n$ = {v[k]['n']}}} \\\\")
    for seg, l in [("segment_tier", "Price-tier markets"), ("segment_body_filled", "Body-type markets")]:
        s = A[seg]; L.append(f"{l}: PCA share / held-out slope & \\multicolumn{{4}}{{c}}{{{s['pca_explained']:.2f} / {100*s['validation']['noexit_fe']['b']:.2f} pp [{100*s['validation']['noexit_fe']['lo']:.2f}, {100*s['validation']['noexit_fe']['hi']:.2f}]}} \\\\")
    L += ["\\bottomrule", "\\end{tabular}"]
    _write("tab07_amli.tex", L)


def run():
    tab04_hedonic(); tab05_hazards(); tab06_competition(); tab07_amli()
    return sorted(p.name for p in T.glob("*.tex"))


if __name__ == "__main__":
    print(run())