#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Step 05 — LaTeX tables.

Converts the step-02/03 CSVs into booktabs tables under
``results/tables/``; the paper inputs these files directly.

Usage:  python scripts/05_make_tables.py
"""
import sys
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from wp11 import config  # noqa: E402

RES = config.REPRODUCED
TAB = config.TABLES


def _f(x, nd=3):
    return f"{x:,.{nd}f}"


def _esc(s: str) -> str:
    return (s.replace("–", "--").replace("~", r"$\sim$")
            .replace("<", "$<$").replace(">", "$>$"))


def t_summary():
    s = pd.read_csv(RES / "summary_stats.csv")
    lines = [r"\begin{tabular}{lrrrrrr}", r"\toprule",
             r"Variable & $n$ & Mean & SD & P10 & Median & P90 \\",
             r"\midrule"]
    for _, r in s.iterrows():
        nd = 0 if r["mean"] > 1000 else 2
        label = (str(r["variable"]).replace("($)", "(\\$)")
                 .replace("(m2)", "(m$^2$)"))
        lines.append(
            f"{label} & {r['n']:,.0f} & {_f(r['mean'], nd)} & "
            f"{_f(r['sd'], nd)} & {_f(r['p10'], nd)} & {_f(r['p50'], nd)} & "
            f"{_f(r['p90'], nd)} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    (TAB / "summary_stats.tex").write_text("\n".join(lines))


def _race_table(fname_csv: str, fname_tex: str):
    r = pd.read_csv(RES / fname_csv)
    lines = [r"\begin{tabular}{lcccc c}", r"\toprule",
             r"Model & MdAPE (\%) & MAPE (\%) & RMSE$_{\ln}$ & "
             r"$R^2_{\ln}$ & Fit (s) \\"]
    for g in ["A. Functional form", "B. Time effects",
              "C. Spatial controls", "D. Estimation method"]:
        sub = r[r["group"] == g]
        if sub.empty:
            continue
        lines.append(r"\midrule")
        lines.append(r"\multicolumn{6}{l}{\itshape " + g + r"}\\")
        best = sub["mdape"].min()
        for _, x in sub.iterrows():
            name = _esc(str(x["name"]))
            md = f"\\textbf{{{_f(x['mdape'], 1)}}}" \
                if x["mdape"] == best else _f(x["mdape"], 1)
            lines.append(
                f"\\quad {name} & {md} & {_f(x['mape'], 1)} & "
                f"{_f(x['rmse_ln'])} & {_f(x['r2_ln'])} & "
                f"{_f(x['seconds'], 0)} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    (TAB / fname_tex).write_text("\n".join(lines))


def t_generalization():
    a = pd.read_csv(RES / "horserace_random.csv")
    b = pd.read_csv(RES / "horserace_temporal.csv")
    m = a.merge(b, on=["name", "group"], suffixes=("_r", "_t"))
    m = m[m["name"].str.match(r"(A5|B4|C4|D)")]
    lines = [r"\begin{tabular}{lcc c cc}", r"\toprule",
             r" & \multicolumn{2}{c}{Random holdout} & &"
             r" \multicolumn{2}{c}{Forward-in-time} \\",
             r"\cmidrule{2-3}\cmidrule{5-6}",
             r"Model & MdAPE (\%) & $R^2_{\ln}$ & & MdAPE (\%) & "
             r"$R^2_{\ln}$ \\", r"\midrule"]
    for _, x in m.iterrows():
        lines.append(
            f"{_esc(str(x['name']))} & {_f(x['mdape_r'], 1)} & "
            f"{_f(x['r2_ln_r'])} & & {_f(x['mdape_t'], 1)} & "
            f"{_f(x['r2_ln_t'])} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    (TAB / "generalization.tex").write_text("\n".join(lines))


def t_boxcox():
    b = pd.read_csv(RES / "boxcox.csv")
    lines = [r"\begin{tabular}{l" + "c" * len(b) + "}", r"\toprule",
             r"$\lambda$ & " + " & ".join(_f(x, 2) for x in b["lambda"])
             + r" \\", r"\midrule",
             r"Profile log-likelihood$^{\dagger}$ & "
             + " & ".join(f"{(x - b['loglik'].max())/1000:,.1f}"
                          for x in b["loglik"]) + r" \\",
             r"\bottomrule", r"\end{tabular}"]
    (TAB / "boxcox.tex").write_text("\n".join(lines))


def t_learning():
    l = pd.read_csv(RES / "learning.csv")
    piv = l.pivot_table(index="n_train", columns="model", values="mdape")
    cols = list(piv.columns)
    lines = [r"\begin{tabular}{r" + "c" * len(cols) + "}", r"\toprule",
             "Training sales & " + " & ".join(_esc(c) for c in cols) + r" \\",
             r"\midrule"]
    for n, row in piv.iterrows():
        lines.append(f"{n:,.0f} & "
                     + " & ".join(_f(v, 1) for v in row) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    (TAB / "learning.tex").write_text("\n".join(lines))


def t_segments():
    s = pd.read_csv(RES / "segments.csv")
    piv = s.pivot_table(index="segment", columns="model", values="mdape")
    ns = s.groupby("segment")["n"].first()
    lines = [r"\begin{tabular}{lrcc}", r"\toprule",
             r"Segment & $n$ (test) & OLS quadratic & Gradient boosting \\",
             r"\midrule"]
    for seg, row in piv.iterrows():
        label = _esc(str(seg)).replace("Class: ", "").replace("_", " ")
        lines.append(f"{label.capitalize()} & {ns[seg]:,.0f} & "
                     f"{_f(row['OLS quadratic'], 1)} & "
                     f"{_f(row['Gradient boosting'], 1)} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    (TAB / "segments.tex").write_text("\n".join(lines))


def main() -> None:
    config.ensure_dirs()
    t_summary()
    _race_table("horserace_random.csv", "horserace_random.tex")
    _race_table("horserace_temporal.csv", "horserace_temporal.tex")
    t_generalization()
    t_boxcox()
    t_learning()
    t_segments()
    made = sorted(p.name for p in TAB.glob("*.tex"))
    print(f"{len(made)} tables written:", ", ".join(made))


if __name__ == "__main__":
    main()