# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""
Shared LaTeX table helpers for the regression scripts.

`significance_star` and `results_to_latex` reproduce byte-for-byte the table
format of the published tables (stargazer-style booktabs tables with HC1
robust standard errors in parentheses).

Tables are written as *fragments* (no ``\\begin{table}`` wrapper): the paper
wraps each ``\\input`` in its own ``table`` environment with caption and label.
"""

from pathlib import Path


def significance_star(pval: float) -> str:
    """Return conventional significance stars for a p-value."""
    if pval < 0.01:
        return "***"
    elif pval < 0.05:
        return "**"
    elif pval < 0.10:
        return "*"
    return ""


def results_to_latex(
    results_list: list,
    model_names: list[str],
    dep_var: str,
    display_vars: list[str] | None = None,
    out_path: Path | None = None,
    note: str = "",
) -> str:
    """
    Produce a stargazer-style LaTeX table from a list of OLS results.

    Parameters
    ----------
    results_list : list of statsmodels RegressionResultsWrapper
    model_names  : column headers for each model
    dep_var      : dependent variable label for the table caption row
    display_vars : subset of variable names to display (None = all)
    out_path     : file to write; None = no write
    note         : optional footnote text
    """
    n_models = len(results_list)

    # Collect the union of variable names across all models
    if display_vars is None:
        all_vars: list[str] = []
        for res in results_list:
            for v in res.params.index:
                if v not in all_vars:
                    all_vars.append(v)
        display_vars = all_vars

    # Build LaTeX
    col_spec = "l" + "c" * n_models
    lines: list[str] = []
    lines.append(r"\small")
    lines.append(r"\begin{tabular}{" + col_spec + "}")
    lines.append(r"\toprule")

    # Header
    header = " & ".join([""] + [f"\\textbf{{{mn}}}" for mn in model_names]) + r" \\"
    lines.append(header)
    lines.append(
        " & ".join(["Dep.\\ var:"] + [f"\\textit{{{dep_var}}}" for _ in model_names])
        + r" \\"
    )
    lines.append(r"\midrule")

    # Coefficients
    for var in display_vars:
        coef_cells = []
        se_cells = []
        for res in results_list:
            if var in res.params.index:
                b = res.params[var]
                se = res.bse[var]
                p = res.pvalues[var]
                star = significance_star(p)
                coef_cells.append(f"{b:.4f}{star}")
                se_cells.append(f"({se:.4f})")
            else:
                coef_cells.append("")
                se_cells.append("")
        # Pretty variable name
        vname = var.replace("_", r"\_")
        lines.append(" & ".join([vname] + coef_cells) + r" \\")
        lines.append(" & ".join([""] + se_cells) + r" \\[4pt]")

    lines.append(r"\midrule")

    # Fit statistics
    for label, accessor in [
        ("Observations", lambda r: f"{int(r.nobs)}"),
        ("R$^2$", lambda r: f"{r.rsquared:.4f}"),
        ("Adj.\\ R$^2$", lambda r: f"{r.rsquared_adj:.4f}"),
    ]:
        cells = [label]
        for res in results_list:
            try:
                cells.append(accessor(res))
            except Exception:
                cells.append("")
        lines.append(" & ".join(cells) + r" \\")

    lines.append(r"\bottomrule")
    lines.append(r"\end{tabular}")

    # Note
    if note:
        lines.append(r"\vspace{4pt}")
        lines.append(r"\parbox{\textwidth}{\footnotesize " + note + "}")

    lines.append(
        r"\parbox{\textwidth}{\footnotesize Robust (HC1) standard errors in "
        r"parentheses. $^{***}p<0.01$; $^{**}p<0.05$; $^{*}p<0.10$.}"
    )

    tex = "\n".join(lines) + "\n"

    if out_path is not None:
        out_path.write_text(tex, encoding="utf-8")
        print(f"  -> saved {out_path}")

    return tex