"""Shared estimation and reporting routines for the editable Quarto template."""

from __future__ import annotations

import hashlib
import json
import os
import subprocess
import tempfile
from importlib.metadata import version
from pathlib import Path
from unittest.mock import patch

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import polars as pl
import pyfixest as pf
from IPython.display import Markdown, display
from pyfixest.did import visualize
from scipy import stats

METHODS = ["Static TWFE", "Dynamic TWFE", "BJS / did2s", "Sun–Abraham"]


def text(value):
    display(Markdown(value))


def table(value):
    display(Markdown(pd.DataFrame(value).to_markdown(index=False, floatfmt=".5g")))


def clean(value):
    if isinstance(value, dict):
        return {str(k): clean(v) for k, v in value.items()}
    if isinstance(value, (list, tuple, np.ndarray)):
        return [clean(x) for x in value]
    if isinstance(value, (np.integer, np.bool_)):
        return value.item()
    if isinstance(value, (float, np.floating)):
        return float(value) if np.isfinite(value) else None
    return value


def save_json(path, value):
    Path(path).write_text(json.dumps(clean(value), indent=2, allow_nan=False) + "\n")


def linear(fit, r, df):
    estimate = float(r @ fit._beta_hat)
    variance = float(r @ fit._vcov @ r)
    if not np.isfinite(variance) or variance < -1e-9:
        raise ValueError("Invalid contrast covariance")
    se = np.sqrt(max(0, variance))
    critical = stats.t.ppf(0.975, df)
    return dict(
        estimate=estimate,
        se=se,
        pvalue=float(2 * stats.t.sf(abs(estimate / se), df)) if se else None,
        lower=estimate - critical * se,
        upper=estimate + critical * se,
    )


def keep_first_stage_singletons(*args, **kwargs):
    # One untreated observation identifies a unit intercept for imputation.
    # did2s does not expose feols' singleton option, so scope it to this call.
    kwargs.setdefault("fixef_rm", "none")
    return pf.feols(*args, **kwargs)


class Replication:
    def __init__(self, config, project):
        self.root = Path(project)
        self.config = (
            dict(
                outcomes=["y"],
                treatments=["d"],
                unit="unit",
                time="time",
                cluster="cluster",
                covariates=[],
                baseline_by_time=[],
                weight=None,
                period_step=1,
                period_values=None,
                display_window=[-10, 10],
                seed=20261008,
                panelview_max_units=100,
                test_mode="classical",
                test_horizon=None,
                post_horizon=None,
                post_weighting="observations",
                sample_query=None,
                reversal_policy="error",
                labels={},
                **{},
            )
            | config
        )
        self.path = self.root / self.config["data"]
        self.output = self.root / self.config["output_dir"]
        self.output.mkdir(parents=True, exist_ok=True)
        if self.path.suffix == ".csv":
            self.df = pl.read_csv(self.path)
        elif self.path.suffix in {".parquet", ".pq"}:
            self.df = pl.read_parquet(self.path)
        elif self.path.suffix == ".dta":
            self.df = pl.from_pandas(
                pd.read_stata(self.path, convert_categoricals=False)
            )
        else:
            raise ValueError("Input must be CSV, Parquet or Stata")
        self.results = dict(
            schema="did-replication-v2",
            config=self.config,
            source=dict(
                file=self.path.name,
                bytes=self.path.stat().st_size,
                sha256=hashlib.file_digest(self.path.open("rb"), "sha256").hexdigest(),
            ),
            versions={
                p: version(p)
                for p in ["pyfixest", "polars", "duckreg", "numpy", "pandas"]
            },
            specifications={},
        )
        self.panels, self.fits = {}, {}

    def label(self, name):
        return self.config["labels"].get(name, name)

    def describe(self):
        c = self.config
        text(
            f"**{c['study']}** — {c['citation']}\n\n{c['description']}\n\n"
            f"Archive specification: `{c['source_specification_id']}`. "
            f"Input: {self.path.stat().st_size / 1e6:.3f} MB."
        )
        rows = [("Y", x, self.label(x)) for x in c["outcomes"]]
        rows += [
            (
                "W",
                x,
                f"{self.label(x)}; values {sorted(self.df[x].unique().to_list())}; cardinality {self.df[x].n_unique()}",
            )
            for x in c["treatments"]
        ]
        rows += [
            ("i", c["unit"], self.label(c["unit"])),
            ("t", c["time"], self.label(c["time"])),
            (
                "X_it",
                ", ".join(
                    c["covariates"] + [x + " × period" for x in c["baseline_by_time"]]
                )
                or "None",
                "Plus unit and period fixed effects",
            ),
            (
                "Cluster",
                c["cluster"] or c["unit"],
                self.label(c["cluster"] or c["unit"]),
            ),
        ]
        table(pd.DataFrame(rows, columns=["Role", "Column", "Definition"]))

    def prepare(self):
        c = self.config
        raw = self.df.to_pandas()
        # Peter's scale is the unweighted SD of the entire prepared source sample.
        scales = {y: float(raw[y].std(ddof=1)) for y in c["outcomes"]}
        if c["sample_query"]:
            raw = raw.query(c["sample_query"]).copy()
        u, t, cl = c["unit"], c["time"], c["cluster"] or c["unit"]
        assert not raw[[u, t, cl]].isna().any().any()
        assert not raw.duplicated([u, t]).any(), "Input must be a unit-period panel"
        if c["period_values"] is not None:
            idx = raw[t].map({v: j + 1 for j, v in enumerate(c["period_values"])})
            assert idx.notna().all()
        else:
            idx = (raw[t] - raw[t].min()) / c["period_step"] + 1
            assert np.allclose(idx, np.round(idx)), "Supply the complete period grid"
        self.period_labels = dict(zip(idx.astype(int), raw[t]))
        samples = []
        for w in c["treatments"]:
            base = pd.DataFrame(
                dict(
                    unit=pd.factorize(raw[u], sort=True)[0] + 1,
                    time=idx.astype(int),
                    d=raw[w],
                    cluster=pd.factorize(raw[cl], sort=True)[0] + 1,
                ),
                index=raw.index,
            )
            assert base.d.notna().all() and set(base.d.unique()) <= {0, 1}
            base = base.sort_values(["unit", "time"])
            base["g"] = (
                base.time.where(base.d.eq(1))
                .groupby(base.unit)
                .transform("min")
                .fillna(0)
                .astype(int)
            )
            base["event_time"] = (base.time - base.g).where(base.g.gt(0))
            base["ever"] = base.g.gt(0).astype(int)
            base["relative_time"] = base.event_time.fillna(-1).astype(int)
            reversals = base.g.gt(0) & base.event_time.ge(0) & base.d.eq(0)
            reversal_count = int(reversals.sum())
            if reversal_count:
                assert c["reversal_policy"] == "drop_post_untreated", (
                    "Nonabsorbing treatment"
                )
                base = base.loc[~reversals].copy()
            always = base.loc[
                base.g.gt(0) & base.g.le(base.groupby("unit").time.transform("min")),
                "unit",
            ].unique()
            base = base.loc[~base.unit.isin(always)].copy()
            xcols = []
            for j, x in enumerate(c["covariates"]):
                name = f"x{j}"
                base[name] = pd.to_numeric(raw.loc[base.index, x])
                xcols.append(name)
            for x in c["baseline_by_time"]:
                values = pd.to_numeric(raw.loc[base.index, x])
                assert values.groupby(base.unit).nunique().max() <= 1, (
                    "Not a unit baseline"
                )
                values = (values - values.mean()) / values.std(ddof=1)
                for period in sorted(base.time.unique())[1:]:
                    name = f"x{len(xcols)}"
                    base[name] = values * base.time.eq(period)
                    xcols.append(name)
            base["weight"] = raw.loc[base.index, c["weight"]] if c["weight"] else 1.0
            for y in c["outcomes"]:
                assert np.isfinite(scales[y]) and scales[y] > 0
                data = base.assign(y=raw.loc[base.index, y] / scales[y])
                data = data.replace([np.inf, -np.inf], np.nan).dropna(
                    subset=["y", "weight", *xcols]
                )
                assert data.weight.gt(0).all()
                data = data.loc[
                    data.g.eq(0) | data.g.gt(data.groupby("unit").time.transform("min"))
                ].copy()
                baseline_units = set(data.loc[data.event_time.eq(-1), "unit"])
                baseline_excluded = data.loc[
                    data.g.gt(0) & ~data.unit.isin(baseline_units), "unit"
                ].nunique()
                data = data.loc[data.g.eq(0) | data.unit.isin(baseline_units)].copy()
                # Remove recursively uninformative unit/period fixed-effect singletons.
                # Recheck pre-treatment support after these removals.
                while True:
                    before = len(data)
                    data = data.loc[
                        data.groupby("unit").unit.transform("size").gt(1)
                        & data.groupby("time").time.transform("size").gt(1)
                    ].copy()
                    data = data.loc[
                        data.g.eq(0)
                        | data.g.gt(data.groupby("unit").time.transform("min"))
                    ].copy()
                    if len(data) == before:
                        break
                recovered = (
                    data.time.where(data.d.eq(1))
                    .groupby(data.unit)
                    .transform("min")
                    .fillna(0)
                    .astype(int)
                )
                assert recovered.equals(data.g), (
                    "Missing outcomes obscure the original adoption date"
                )
                assert data.g.eq(0).any() and data.loc[data.g.gt(0), "g"].nunique() >= 2
                assert data.groupby("unit").size().min() >= 2
                data = data.reset_index(drop=True)
                key = f"{y}__{w}"
                self.panels[key] = (data, xcols)
                units = data[["unit", "g"]].drop_duplicates()
                sample = dict(
                    outcome=y,
                    treatment=w,
                    source_rows=self.df.height,
                    source_sd=scales[y],
                    rows=len(data),
                    units=len(units),
                    clusters=data.cluster.nunique(),
                    treated_units=int(units.g.gt(0).sum()),
                    never_treated_units=int(units.g.eq(0).sum()),
                    treated_cohorts=units.loc[units.g.gt(0), "g"].nunique(),
                    initially_treated_units_removed=len(always),
                    post_adoption_untreated_rows_removed=reversal_count,
                    units_without_minus_one_removed=baseline_excluded,
                    additional_rows_removed=len(base) - len(data),
                )
                self.results["specifications"][key] = dict(sample=sample)
                samples.append(sample)
        return pd.DataFrame(samples)

    def panelview(self):
        for key, (data, _) in self.panels.items():
            n = min(self.config["panelview_max_units"], data.unit.nunique())
            chosen = np.random.default_rng(self.config["seed"]).choice(
                sorted(data.unit.unique()), n, replace=False
            )
            fig, ax = plt.subplots(figsize=(9, 4))
            plot_quilt = visualize._plot_panelview

            def numeric_quilt(treatment_quilt, **kwargs):
                # Boolean pivots with missing cells otherwise become object arrays.
                return plot_quilt(treatment_quilt.astype(float), **kwargs)

            with patch("pyfixest.did.visualize._plot_panelview", numeric_quilt):
                pf.panelview(
                    data.loc[data.unit.isin(chosen)],
                    unit="unit",
                    time="time",
                    treat="d",
                    sort_by_timing=True,
                    noticks=True,
                    ax=ax,
                    title="Treatment adoption",
                )
            periods = sorted(data.time.unique())
            positions = np.unique(
                np.linspace(0, len(periods) - 1, min(7, len(periods)), dtype=int)
            )
            ax.set_xticks(
                positions, [str(self.period_labels[periods[k]]) for k in positions]
            )
            self._show(fig)
            text(
                f"{key}: {n} of {data.unit.nunique()} units shown (seed {self.config['seed']}). "
                "All eligible units enter estimation. Purple: untreated; yellow: treated; white: unobserved."
            )

    def estimate(self):
        for key, (data, xs) in self.panels.items():
            print(f"Fitting {key}: {len(data):,} rows", flush=True)
            extra = (" + " + " + ".join(xs)) if xs else ""
            args = dict(
                data=data,
                vcov={"CRV1": "cluster"},
                weights="weight" if self.config["weight"] else None,
            )
            static = pf.feols("y ~ d" + extra + " | unit + time", **args)
            dynamic = pf.feols(
                "y ~ i(relative_time, ever, ref=-1)" + extra + " | unit + time", **args
            )
            # Explicit second-stage indicators: never-treated rows are zero in every column.
            events = sorted(data.event_time.dropna().astype(int).unique())
            columns = {
                e: f"es_{'m' + str(-e) if e < 0 else 'p' + str(e)}" for e in events
            }
            indicators = pd.DataFrame(
                {
                    name: data.event_time.eq(e).astype(float)
                    for e, name in columns.items()
                }
            )
            with patch("pyfixest.did.did2s.feols", keep_first_stage_singletons):
                did = pf.did2s(
                    pd.concat([data, indicators], axis=1),
                    yname="y",
                    first_stage="~ "
                    + (" + ".join(xs) if xs else "0")
                    + " | unit + time",
                    second_stage="~ " + " + ".join(columns.values()),
                    treatment="d",
                    cluster="cluster",
                    weights=args["weights"],
                )
            sa = pf.event_study(
                data=data,
                yname="y",
                idname="unit",
                tname="time",
                gname="g",
                cluster="cluster",
                estimator="saturated",
                att=False,
            )
            if xs or self.config["weight"]:
                lhs, fe = sa._fml.split("|")
                sa = pf.feols(
                    lhs + extra + " | " + fe,
                    data=sa._data,
                    vcov={"CRV1": "cluster"},
                    weights=args["weights"],
                )
            assert all(m._N == len(data) for m in [static, dynamic, did, sa])
            fits = dict(zip(METHODS, [static, dynamic, did, sa]))
            self.fits[key] = fits
            self._summarize(key, data, fits, columns)

    def _summarize(self, key, data, fits, columns):
        df = data.cluster.nunique() - 1
        events = sorted(data.event_time.dropna().astype(int).unique())
        mass = data.groupby(["g", "event_time"]).weight.sum()
        sa = fits["Sun–Abraham"]
        parts = (
            pd.Series(sa.coef().index)
            .str.extract(
                r"^rel_time::(?P<e>-?\d+(?:\.\d+)?):first_treated_period::(?P<g>\d+)$"
            )
            .astype(float)
        )
        sa_rows = {}
        for e in sorted(parts.e.dropna().unique()):
            hit = parts.e.eq(e)
            r = np.zeros(len(parts))
            r[hit] = [mass.loc[(int(g), e)] for g in parts.loc[hit, "g"]]
            sa_rows[int(e)] = r / r.sum()
        package_aggregation_gap = None
        if not (
            self.config["covariates"]
            or self.config["baseline_by_time"]
            or self.config["weight"]
        ):
            official = sa.aggregate(agg="period", weighting="shares").astype(float)
            expected = np.array(
                [linear(sa, r, df)["estimate"] for r in sa_rows.values()]
            )
            package_aggregation_gap = float(
                np.max(np.abs(expected - official["Estimate"].to_numpy()))
            )
            assert package_aggregation_gap < 1e-7, (
                "Saturated aggregation needs a support review"
            )
        # Renormalize over estimated cohort cells; omitted columns get no target mass.
        dynamic = fits["Dynamic TWFE"]
        dyn_events = (
            pd.Series(dynamic.coef().index)
            .str.extract(r"^relative_time::(-?\d+(?:\.\d+)?)", expand=False)
            .astype(float)
        )
        matrices = {
            "Dynamic TWFE": {
                e: dyn_events.eq(e).astype(float).to_numpy()
                for e in events
                if dyn_events.eq(e).any()
            },
            "BJS / did2s": {
                e: (fits["BJS / did2s"].coef().index == name).astype(float)
                for e, name in columns.items()
            },
            "Sun–Abraham": sa_rows,
        }
        post = data.loc[data.d.eq(1)].groupby("event_time").weight.sum()
        if self.config["post_horizon"] is not None:
            post = post.loc[post.index <= self.config["post_horizon"]]
        common = sorted(
            set(post.index).intersection(*(set(m) for m in matrices.values()))
        )
        assert common
        weights = post.reindex(common).to_numpy(copy=True)
        if self.config["post_weighting"] == "equal_periods":
            weights = np.ones_like(weights)
        else:
            assert self.config["post_weighting"] == "observations"
        weights /= weights.sum()
        pooled = [
            dict(
                method="Static TWFE",
                **linear(
                    fits["Static TWFE"],
                    (fits["Static TWFE"].coef().index == "d").astype(float),
                    df,
                ),
            )
        ]
        paths = {}
        for method, matrix in matrices.items():
            model = fits[method]
            r = sum(w * matrix[e] for e, w in zip(common, weights))
            pooled.append(dict(method=method, **linear(model, r, df)))
            paths[method] = [
                dict(event_time=e, **linear(model, r, df)) for e, r in matrix.items()
            ]
        # Independent imputation check of every post-treatment did2s coefficient.
        xs = self.panels[key][1]
        first_formula = "y ~ " + (" + ".join(xs) if xs else "0") + " | unit + time"
        first = pf.feols(
            first_formula,
            data=data.loc[data.d.eq(0)],
            weights="weight" if self.config["weight"] else None,
            fixef_tol=1e-10,
            fixef_rm="none",
        )
        first.fixef()
        residual = data.y.to_numpy() - first.predict(newdata=data)
        mass_by_event = data.groupby("event_time").weight.sum()
        weighted = (
            pd.Series(residual * data.weight.to_numpy()).groupby(data.event_time).sum()
        )
        bjs = weighted / mass_by_event
        check = max(
            abs(row["estimate"] - bjs.loc[row["event_time"]])
            for row in paths["BJS / did2s"]
            if row["event_time"] >= 0
        )
        assert check < 1e-5, f"did2s/imputation discrepancy: {check}"
        output = self.results["specifications"][key]
        output.update(
            estimates=pooled,
            event_studies=paths,
            target=dict(
                event_times=common,
                weights=weights.tolist(),
                weighting=self.config["post_weighting"],
            ),
            checks=dict(
                did2s_vs_direct_imputation_max_post_difference=check,
                package_aggregate_vs_estimable_cell_weights=package_aggregation_gap,
            ),
            inference=dict(
                standard_errors="Cluster-robust; did2s uses two-stage covariance",
                confidence_intervals="Pointwise t, clusters minus one degrees of freedom",
            ),
        )
        pd.DataFrame(pooled).to_csv(self.output / f"{key}-estimates.csv", index=False)

    def estimate_table(self):
        for key, result in self.results["specifications"].items():
            text(f"**{key}**")
            table(result["estimates"])
            target = result["target"]
            text(
                f"Post-treatment event times {min(target['event_times']):g}–{max(target['event_times']):g}; "
                f"{target['weighting'].replace('_', ' ')} weights. All estimates are in source-outcome SD units."
            )

    def event_plots(self):
        for key, result in self.results["specifications"].items():
            static = result["estimates"][0]
            for method, path in result["event_studies"].items():
                frame = pd.DataFrame(path).sort_values("event_time")
                lo, hi = self.config["display_window"]
                frame = frame.loc[frame.event_time.between(lo, hi)]
                fig, ax = plt.subplots(figsize=(8, 3.6))
                ax.errorbar(
                    frame.event_time,
                    frame.estimate,
                    yerr=[frame.estimate - frame.lower, frame.upper - frame.estimate],
                    fmt="o",
                    ms=3,
                    capsize=2,
                    color="#147d87",
                )
                if method == "Dynamic TWFE":
                    ax.scatter(
                        [-1], [0], facecolor="white", edgecolor="#147d87", zorder=4
                    )
                    ax.axhline(static["estimate"], color="#a15236", label="Static TWFE")
                    ax.axhspan(
                        static["lower"], static["upper"], alpha=0.08, color="#a15236"
                    )
                    ax.legend()
                elif method == "Sun–Abraham":
                    ax.scatter(
                        [-1], [0], facecolor="white", edgecolor="#147d87", zorder=4
                    )
                ax.axhline(0, color="gray", lw=0.7)
                ax.axvline(-0.5, color="gray", ls="--", lw=0.7)
                ax.set(
                    title=method,
                    xlabel="Periods relative to adoption",
                    ylabel="Outcome SD units",
                )
                self._show(fig)

    def f_tests(self):
        for key, (data, xs) in self.panels.items():
            outfile = self.output / f"{key}-f-test.json"
            with tempfile.TemporaryDirectory() as td:
                input_path = Path(td) / "panel.csv"
                data.to_csv(input_path, index=False)
                conf = dict(
                    reference_repo=str(self.root / "TestingInEventStudies"),
                    data=str(input_path),
                    unit="unit",
                    time="time",
                    outcome="y",
                    treatment="d",
                    cluster="cluster",
                    covariates=xs,
                    weight="weight" if self.config["weight"] else None,
                    test_mode=self.config["test_mode"],
                    test_horizon=self.config["test_horizon"],
                    output=str(outfile),
                )
                result = subprocess.run(
                    [
                        str(self.root / "TestingInEventStudies/.venv/bin/python"),
                        str(Path(__file__).with_name("reference_f_test.py")),
                        json.dumps(conf),
                    ],
                    env=dict(
                        os.environ,
                        MPLBACKEND="Agg",
                        OPENBLAS_NUM_THREADS="1",
                        OMP_NUM_THREADS="1",
                        NUMBA_NUM_THREADS="2",
                        PYTHONDONTWRITEBYTECODE="1",
                    ),
                    capture_output=True,
                    text=True,
                )
                if result.returncode:
                    raise RuntimeError(result.stderr)
            test = json.loads(outfile.read_text())
            sa = self.fits[key]["Sun–Abraham"]
            sa_residual = np.asarray(sa._u_hat).ravel()
            sa_rss = float(sa_residual @ sa_residual)
            gap = abs(sa_rss - test["rss_unrestricted"]) / max(1, sa_rss)
            assert gap < 1e-7, f"Saturated / deviations fit mismatch: {gap}"
            self.results["specifications"][key]["checks"][
                "saturated_vs_deviations_relative_rss_difference"
            ] = gap
            self.results["specifications"][key]["f_test"] = test

    def f_report(self):
        for key, result in self.results["specifications"].items():
            test = result["f_test"]
            text(
                f"**{key}: F({test['df_num']}, {test['df_den']}) = {test['F']:.4f}; p = {test['pvalue']:.4g}.** "
                f"{test['inference']}. {test['design_columns']} design columns."
            )
            frame = pd.DataFrame(test["deviations"])
            lo, hi = self.config["display_window"]
            frame = frame.loc[frame.event_time.between(lo, hi)]
            fig, ax = plt.subplots(figsize=(8, 3.7))
            for cohort, subset in frame.groupby("cohort"):
                subset = subset.sort_values("event_time")
                ax.plot(subset.event_time, subset.estimate, lw=0.9, alpha=0.65)
            ax.axhline(0, color="black", lw=0.8)
            ax.axvline(-0.5, color="gray", ls="--")
            ax.set(
                title="Cohort deviations from the reference event-time path",
                xlabel="Periods relative to adoption",
                ylabel="Deviation (outcome SD units)",
            )
            self._show(fig)
            text(
                "Each line is a cohort deviation. The joint test restricts positive-event-time deviations to zero; event time 0 remains unrestricted."
            )

    def code_hashes(self):
        return {
            name: hashlib.sha256(
                Path(__file__).with_name(name).read_bytes()
            ).hexdigest()
            for name in ["replication_engine.py", "reference_f_test.py"]
        }

    def load_results(self):
        cached = json.loads((self.output / "estimates.json").read_text())
        assert cached["config"] == self.config, "Cached configuration changed"
        assert cached["source"] == self.results["source"], "Cached source changed"
        assert cached["versions"] == self.results["versions"], (
            "Cached environment changed"
        )
        assert cached["code_sha256"] == self.code_hashes(), (
            "Cached implementation changed"
        )
        self.results = cached

    def export(self):
        self.results["code_sha256"] = self.code_hashes()
        save_json(self.output / "estimates.json", self.results)
        return self.output / "estimates.json"

    @staticmethod
    def _show(fig):
        fig.tight_layout()
        plt.show()
        plt.close(fig)