from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor
from dataclasses import asdict, dataclass
from hashlib import sha256
from pathlib import Path
import argparse
import json
import math
import os
import platform
import sys

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

from common_models import (
    MODEL_NAMES,
    UNIVERSE_SIZES,
    acceptable_indices,
    causal_ewma_adjust,
    fit_models,
    model_quantile_es,
    observation_log_losses,
    score_weights,
)


HERE = Path(__file__).resolve().parent
PROJECT = HERE.parent
DEFAULT_OUTPUT = PROJECT / "results" / "study1"
PROTOCOL = PROJECT / "08_FROZEN_CONFIRMATORY_PROTOCOL.md"


@dataclass(frozen=True)
class Config:
    seed: int = 2026072602
    reps_per_dgp: int = 400
    training_n: int = 1000
    validation_n: int = 500
    test_n: int = 250
    evidence_sizes: tuple[int, ...] = (20, 60, 120, 250, 500)
    alpha: float = 0.05
    negligible_margin: float = 0.02


DGP_NAMES = (
    "stable_gaussian",
    "stable_student_t",
    "stable_skew_normal",
    "stable_mixture",
    "garch_scale_adjusted",
    "single_break",
    "markov_switching",
    "smooth_drift",
)
STABLE_DGPS = DGP_NAMES[:4]


def _standardized_skew(
    rng: np.random.Generator, size: int, alpha: float = 7
) -> np.ndarray:
    raw = stats.skewnorm.rvs(alpha, size=size, random_state=rng)
    mean, variance = stats.skewnorm.stats(alpha, moments="mv")
    return (raw - float(mean)) / np.sqrt(float(variance))


def _stable_mixture(rng: np.random.Generator, size: int) -> np.ndarray:
    weight = 0.15
    left_mean, left_sd = -3.2, 0.30
    right_mean, right_sd = -weight * left_mean / (1 - weight), 0.70
    variance = (
        weight * (left_sd**2 + left_mean**2)
        + (1 - weight) * (right_sd**2 + right_mean**2)
    )
    choose_left = rng.uniform(size=size) < weight
    raw = np.where(
        choose_left,
        rng.normal(left_mean, left_sd, size),
        rng.normal(right_mean, right_sd, size),
    )
    return raw / np.sqrt(variance)


def draw_path(
    dgp: str, rng: np.random.Generator, config: Config
) -> np.ndarray:
    total = config.training_n + 2 * config.validation_n + config.test_n
    if dgp == "stable_gaussian":
        values = rng.normal(size=total)
    elif dgp == "stable_student_t":
        values = rng.standard_t(5, size=total) / np.sqrt(5 / 3)
    elif dgp == "stable_skew_normal":
        values = _standardized_skew(rng, total)
    elif dgp == "stable_mixture":
        values = _stable_mixture(rng, total)
    elif dgp == "garch_scale_adjusted":
        shocks = rng.normal(size=total)
        variance = np.ones(total)
        raw = np.zeros(total)
        for index in range(1, total):
            variance[index] = (
                0.03 + 0.10 * raw[index - 1] ** 2 + 0.87 * variance[index - 1]
            )
            raw[index] = np.sqrt(variance[index]) * shocks[index]
        values, _ = causal_ewma_adjust(raw)
    elif dgp == "single_break":
        break_at = config.training_n
        values = np.r_[
            rng.normal(size=break_at),
            rng.standard_t(4, size=total - break_at) / np.sqrt(2),
        ]
    elif dgp == "markov_switching":
        state = np.zeros(total, dtype=int)
        for index in range(1, total):
            state[index] = (
                state[index - 1]
                if rng.uniform() < 0.96
                else 1 - state[index - 1]
            )
        normal = rng.normal(size=total)
        skewed = _standardized_skew(rng, total, alpha=-7)
        values = np.where(state == 0, normal, skewed)
    elif dgp == "smooth_drift":
        tail_probability = np.linspace(0.01, 0.20, total)
        tail = rng.normal(-2.8, 0.45, total)
        center = rng.normal(0.08, 0.90, total)
        values = np.where(rng.uniform(size=total) < tail_probability, tail, center)
    else:
        raise ValueError(dgp)
    return np.asarray(values, dtype=float)


def evaluate_replication(
    task: tuple[str, int, Config]
) -> list[dict]:
    dgp, replication, config = task
    dgp_index = DGP_NAMES.index(dgp)
    rng = np.random.default_rng(
        np.random.SeedSequence([config.seed, dgp_index, replication])
    )
    path = draw_path(dgp, rng, config)
    train_end = config.training_n
    first_end = train_end + config.validation_n
    second_end = first_end + config.validation_n
    training = path[:train_end]
    validation_a = path[train_end:first_end]
    validation_b = path[first_end:second_end]
    test = path[second_end:]
    rows: list[dict] = []
    try:
        models = fit_models(
            training, config.seed + dgp_index * 1_000_000 + replication
        )
        losses_a = observation_log_losses(models, validation_a)
        losses_b = observation_log_losses(models, validation_b)
        losses_test = observation_log_losses(models, test)
        decisions = np.asarray(
            [model_quantile_es(model, training) for model in models]
        )
        scale = max(float(np.std(training, ddof=1)), 1e-10)
        for universe, model_count in UNIVERSE_SIZES.items():
            names = MODEL_NAMES[:model_count]
            test_means = losses_test[:, :model_count].mean(axis=0)
            test_order = np.argsort(test_means)
            test_best = int(test_order[0])
            second = int(test_order[1]) if model_count > 1 else test_best
            test_gap = float(test_means[second] - test_means[test_best])
            score_variance = float(
                np.var(
                    losses_test[:, second] - losses_test[:, test_best],
                    ddof=1,
                )
            )
            for evidence_n in config.evidence_sizes:
                current_a = losses_a[:evidence_n, :model_count]
                current_b = losses_b[:evidence_n, :model_count]
                accepted, winner_a, lower_bounds, standard_errors = (
                    acceptable_indices(
                        current_a,
                        alpha=config.alpha,
                        negligible_margin=config.negligible_margin,
                    )
                )
                _, winner_b, _, _ = acceptable_indices(
                    current_b,
                    alpha=config.alpha,
                    negligible_margin=config.negligible_margin,
                )
                weights = score_weights(current_a[:, accepted])
                effective_count = float(1 / np.sum(weights**2))
                var_values = decisions[accepted, 0] / scale
                es_values = decisions[accepted, 1] / scale
                selected_test_loss = float(test_means[winner_a])
                rows.append(
                    {
                        "dgp": dgp,
                        "rep": replication,
                        "universe": universe,
                        "model_count": model_count,
                        "evidence_n": evidence_n,
                        "fit_failed": False,
                        "best_validation_score": float(
                            np.min(current_a.mean(axis=0))
                        ),
                        "winner": names[winner_a],
                        "winner_reversal": winner_a != winner_b,
                        "acceptable_count": len(accepted),
                        "effective_count": effective_count,
                        "var_diameter": float(np.max(var_values) - np.min(var_values)),
                        "es_diameter": float(np.max(es_values) - np.min(es_values)),
                        "sealed_test_regret": selected_test_loss
                        - float(np.min(test_means)),
                        "minimum_test_score_gap": test_gap,
                        "score_difference_variance": score_variance,
                        "maximum_acceptable_lower_bound": float(
                            np.nanmax(lower_bounds[accepted])
                        ),
                        "maximum_acceptable_se": float(
                            np.nanmax(standard_errors[accepted])
                        ),
                    }
                )
    except Exception as error:
        for universe, model_count in UNIVERSE_SIZES.items():
            for evidence_n in config.evidence_sizes:
                rows.append(
                    {
                        "dgp": dgp,
                        "rep": replication,
                        "universe": universe,
                        "model_count": model_count,
                        "evidence_n": evidence_n,
                        "fit_failed": True,
                        "failure_type": type(error).__name__,
                        "failure_message": str(error)[:240],
                    }
                )
    return rows


def run_experiment(config: Config, workers: int) -> pd.DataFrame:
    tasks = [
        (dgp, replication, config)
        for dgp in DGP_NAMES
        for replication in range(config.reps_per_dgp)
    ]
    rows: list[dict] = []
    if workers == 1:
        for task in tasks:
            rows.extend(evaluate_replication(task))
    else:
        with ProcessPoolExecutor(max_workers=workers) as executor:
            for result in executor.map(evaluate_replication, tasks, chunksize=2):
                rows.extend(result)
    return pd.DataFrame(rows)


def contrast_frame(
    data: pd.DataFrame, outcome: str, dgps: tuple[str, ...]
) -> pd.DataFrame:
    subset = data[
        data.dgp.isin(dgps)
        & data.universe.isin(["U5", "U9"])
        & data.evidence_n.isin([20, 500])
        & (~data.fit_failed)
    ]
    pivot = subset.pivot_table(
        index=["dgp", "rep"],
        columns=["universe", "evidence_n"],
        values=outcome,
        aggfunc="first",
    )
    required = [("U5", 20), ("U9", 20), ("U5", 500), ("U9", 500)]
    pivot = pivot.dropna(subset=required)
    pivot["contrast"] = (
        pivot[("U9", 20)]
        - pivot[("U5", 20)]
        - pivot[("U9", 500)]
        + pivot[("U5", 500)]
    )
    return pivot.reset_index()[["dgp", "rep", "contrast"]]


def summarize_contrast(values: np.ndarray) -> dict:
    values = np.asarray(values, dtype=float)
    mean = float(np.mean(values))
    se = float(stats.sem(values))
    critical = float(stats.t.ppf(0.975, len(values) - 1))
    test = stats.ttest_1samp(values, 0)
    return {
        "n": len(values),
        "mean": mean,
        "se": se,
        "ci_low": mean - critical * se,
        "ci_high": mean + critical * se,
        "t": float(test.statistic),
        "p_two_sided": float(test.pvalue),
    }


def analyze(data: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    contrast_rows: list[dict] = []
    for outcome in ("var_diameter", "acceptable_count", "es_diameter"):
        frame = contrast_frame(data, outcome, DGP_NAMES)
        stable = frame[frame.dgp.isin(STABLE_DGPS)]
        stable_by_rep = stable.groupby("rep").contrast.mean().to_numpy()
        contrast_rows.append(
            {
                "outcome": outcome,
                "scope": "stable_dgps_average",
                **summarize_contrast(stable_by_rep),
            }
        )
        for dgp, group in frame.groupby("dgp"):
            contrast_rows.append(
                {
                    "outcome": outcome,
                    "scope": dgp,
                    **summarize_contrast(group.contrast.to_numpy()),
                }
            )
    contrasts = pd.DataFrame(contrast_rows)

    failure_rate = float(
        data.groupby(["dgp", "rep"]).fit_failed.max().mean()
    )
    nested = data[
        data.universe.isin(["U5", "U9"]) & (~data.fit_failed)
    ].pivot_table(
        index=["dgp", "rep", "evidence_n"],
        columns="universe",
        values="best_validation_score",
        aggfunc="first",
    )
    nested_violation = float((nested.U9 > nested.U5 + 1e-10).mean())
    primary_var = contrasts[
        (contrasts.outcome == "var_diameter")
        & (contrasts.scope == "stable_dgps_average")
    ].iloc[0]
    primary_set = contrasts[
        (contrasts.outcome == "acceptable_count")
        & (contrasts.scope == "stable_dgps_average")
    ].iloc[0]
    gates = pd.DataFrame(
        [
            {
                "gate": "nested_best_score",
                "observed": nested_violation,
                "rule": "equals 0",
                "passed": nested_violation == 0,
            },
            {
                "gate": "positive_var_interaction",
                "observed": primary_var["ci_low"],
                "rule": "95% CI lower bound > 0",
                "passed": primary_var["ci_low"] > 0,
            },
            {
                "gate": "positive_set_interaction",
                "observed": primary_set["ci_low"],
                "rule": "95% CI lower bound > 0",
                "passed": primary_set["ci_low"] > 0,
            },
            {
                "gate": "fit_failure_rate",
                "observed": failure_rate,
                "rule": "< 0.05",
                "passed": failure_rate < 0.05,
            },
        ]
    )
    gates = pd.concat(
        [
            gates,
            pd.DataFrame(
                [
                    {
                        "gate": "all_study1_gates",
                        "observed": float(gates.passed.all()),
                        "rule": "equals 1",
                        "passed": bool(gates.passed.all()),
                    }
                ]
            ),
        ],
        ignore_index=True,
    )
    return contrasts, gates


def write_summary(
    output: Path,
    data: pd.DataFrame,
    contrasts: pd.DataFrame,
    gates: pd.DataFrame,
) -> None:
    primary = contrasts[
        contrasts.scope == "stable_dgps_average"
    ].set_index("outcome")
    lines = [
        "# Study 1 Confirmatory Results",
        "",
        f"- Valid cell rows: {(~data.fit_failed).sum():,}",
        f"- Failed cell rows: {data.fit_failed.sum():,}",
        f"- Replication-level failure rate: "
        f"{data.groupby(['dgp', 'rep']).fit_failed.max().mean():.3%}",
        "",
        "## Co-primary interactions",
        "",
        "| Outcome | Estimate | 95% CI | p (two-sided) |",
        "|---|---:|---:|---:|",
    ]
    for outcome in ("var_diameter", "acceptable_count"):
        row = primary.loc[outcome]
        lines.append(
            f"| {outcome} | {row['mean']:.4f} | "
            f"[{row['ci_low']:.4f}, {row['ci_high']:.4f}] | "
            f"{row['p_two_sided']:.4g} |"
        )
    lines.extend(
        [
            "",
            "The interaction is "
            + ("confirmed." if gates.iloc[-1].passed else "not confirmed under the frozen gates."),
            "",
            "## Gate table",
            "",
            gates.to_markdown(index=False),
            "",
            "All DGP-specific estimates are retained in `study1_contrasts.csv`.",
        ]
    )
    (output / "study1_results.md").write_text(
        "\n".join(lines) + "\n", encoding="utf-8"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=min(4, os.cpu_count() or 1))
    parser.add_argument("--reps", type=int, default=400)
    parser.add_argument("--seed", type=int, default=2026072602)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.workers < 1 or args.reps < 1:
        raise ValueError("workers and reps must be positive")
    config = Config(seed=args.seed, reps_per_dgp=args.reps)
    args.output.mkdir(parents=True, exist_ok=True)
    data = run_experiment(config, args.workers)
    contrasts, gates = analyze(data)
    data.to_csv(args.output / "study1_cells.csv", index=False)
    contrasts.to_csv(args.output / "study1_contrasts.csv", index=False)
    gates.to_csv(args.output / "study1_gates.csv", index=False)
    write_summary(args.output, data, contrasts, gates)
    manifest = {
        "config": asdict(config),
        "dgp_names": DGP_NAMES,
        "stable_dgps": STABLE_DGPS,
        "model_names": MODEL_NAMES,
        "universe_sizes": UNIVERSE_SIZES,
        "script_sha256": sha256(Path(__file__).read_bytes()).hexdigest(),
        "common_models_sha256": sha256(
            (HERE / "common_models.py").read_bytes()
        ).hexdigest(),
        "protocol_sha256": sha256(PROTOCOL.read_bytes()).hexdigest(),
        "python": sys.version,
        "platform": platform.platform(),
    }
    (args.output / "study1_manifest.json").write_text(
        json.dumps(manifest, indent=2), encoding="utf-8"
    )


if __name__ == "__main__":
    main()
