from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass
from pathlib import Path
import argparse

import numpy as np
import pandas as pd
from scipy import stats

from common_models import (
    UNIVERSE_SIZES,
    fit_models,
    model_quantile_es,
    newey_west_mean_se,
    observation_log_losses,
)
from run_study1 import Config as PrimaryConfig
from run_study1 import STABLE_DGPS, draw_path


PROJECT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT = PROJECT / "results" / "study1_sensitivity"
MARGINS = (0.0, 0.01, 0.02, 0.04)
RULES = (
    "same_sample_family_bonferroni",
    "same_sample_fixed_u9_critical",
    "same_sample_unadjusted",
    "split_reference_family_bonferroni",
)


@dataclass(frozen=True)
class Config:
    seed: int = 2026072602
    reps: int = 400
    training_n: int = 1000
    validation_n: int = 500
    test_n: int = 250
    alpha: float = 0.05


def accepted_indices(losses, winner, margin, alpha, critical_count):
    critical = stats.norm.ppf(1 - alpha / max(critical_count, 1))
    accepted = np.zeros(losses.shape[1], dtype=bool)
    accepted[winner] = True
    for index in range(losses.shape[1]):
        if index == winner:
            continue
        differences = losses[:, index] - losses[:, winner]
        lower = float(
            np.mean(differences)
            - critical * newey_west_mean_se(differences)
        )
        accepted[index] = lower <= margin
    return np.flatnonzero(accepted)


def one_rep(task):
    dgp, rep, cfg = task
    primary = PrimaryConfig(
        seed=cfg.seed,
        reps_per_dgp=cfg.reps,
        training_n=cfg.training_n,
        validation_n=cfg.validation_n,
        test_n=cfg.test_n,
    )
    dgp_index = STABLE_DGPS.index(dgp)
    rng = np.random.default_rng(
        np.random.SeedSequence([cfg.seed, dgp_index, rep])
    )
    path = draw_path(dgp, rng, primary)
    t0 = cfg.training_n
    v1 = t0 + cfg.validation_n
    v2 = v1 + cfg.validation_n
    training = path[:t0]
    validation_a = path[t0:v1]
    validation_b = path[v1:v2]
    models = fit_models(training, cfg.seed + dgp_index * 1_000_000 + rep)
    loss_a = observation_log_losses(models, validation_a)
    loss_b = observation_log_losses(models, validation_b)
    decisions = np.asarray([model_quantile_es(m, training)[0] for m in models])
    scale = max(float(np.std(training, ddof=1)), 1e-10)
    decisions = decisions / scale
    rows = []
    for universe in ("U5", "U9"):
        k = UNIVERSE_SIZES[universe]
        for evidence_n in (20, 500):
            a = loss_a[:evidence_n, :k]
            b = loss_b[:evidence_n, :k]
            winner_a = int(np.argmin(a.mean(axis=0)))
            winner_b = int(np.argmin(b.mean(axis=0)))
            for margin in MARGINS:
                for rule in RULES:
                    if rule == "same_sample_family_bonferroni":
                        winner = winner_a
                        critical_count = k - 1
                    elif rule == "same_sample_fixed_u9_critical":
                        winner = winner_a
                        critical_count = 8
                    elif rule == "same_sample_unadjusted":
                        winner = winner_a
                        critical_count = 1
                    else:
                        winner = winner_b
                        critical_count = k - 1
                    accepted = accepted_indices(
                        a, winner, margin, cfg.alpha, critical_count
                    )
                    vals = decisions[accepted]
                    rows.append(
                        {
                            "dgp": dgp,
                            "rep": rep,
                            "universe": universe,
                            "evidence_n": evidence_n,
                            "margin": margin,
                            "rule": rule,
                            "acceptable_count": len(accepted),
                            "var_diameter": float(vals.max() - vals.min()),
                        }
                    )
    return rows


def summarize(cells):
    rows = []
    for (margin, rule, outcome), group in cells.melt(
        id_vars=["dgp", "rep", "universe", "evidence_n", "margin", "rule"],
        value_vars=["var_diameter", "acceptable_count"],
        var_name="outcome",
        value_name="value",
    ).groupby(["margin", "rule", "outcome"]):
        pivot = group.pivot_table(
            index=["dgp", "rep"],
            columns=["universe", "evidence_n"],
            values="value",
            aggfunc="first",
        )
        pivot["contrast"] = (
            pivot[("U9", 20)]
            - pivot[("U5", 20)]
            - pivot[("U9", 500)]
            + pivot[("U5", 500)]
        )
        per_rep = pivot.reset_index().groupby("rep").contrast.mean()
        mean = float(per_rep.mean())
        se = float(stats.sem(per_rep))
        crit = float(stats.t.ppf(0.975, len(per_rep) - 1))
        rows.append(
            {
                "margin": margin,
                "rule": rule,
                "outcome": outcome,
                "replications": len(per_rep),
                "interaction": mean,
                "se": se,
                "ci_low": mean - crit * se,
                "ci_high": mean + crit * se,
            }
        )
    return pd.DataFrame(rows)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--reps", type=int, default=400)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()
    cfg = Config(reps=args.reps)
    tasks = [(dgp, rep, cfg) for dgp in STABLE_DGPS for rep in range(cfg.reps)]
    rows = []
    if args.workers == 1:
        for task in tasks:
            rows.extend(one_rep(task))
    else:
        with ProcessPoolExecutor(max_workers=args.workers) as pool:
            for result in pool.map(one_rep, tasks):
                rows.extend(result)
    cells = pd.DataFrame(rows)
    summary = summarize(cells)
    args.output.mkdir(parents=True, exist_ok=True)
    cells.to_csv(args.output / "sensitivity_cells.csv", index=False)
    summary.to_csv(args.output / "sensitivity_summary.csv", index=False)
    lines = [
        "# Post-outcome Study 1 Sensitivity",
        "",
        "These analyses were specified after the frozen primary result and do not replace it.",
        "",
        "| Margin | Rule | Outcome | Interaction | 95% CI |",
        "|---:|---|---|---:|---:|",
    ]
    for row in summary.itertuples():
        lines.append(
            f"| {row.margin:.2f} | {row.rule} | {row.outcome} | "
            f"{row.interaction:.4f} | [{row.ci_low:.4f}, {row.ci_high:.4f}] |"
        )
    (args.output / "sensitivity_results.md").write_text(
        "\n".join(lines) + "\n", encoding="utf-8"
    )
    print(args.output)


if __name__ == "__main__":
    main()
