from __future__ import annotations

import json
import math
import os
from pathlib import Path
from typing import Any

os.environ.setdefault("MPLCONFIGDIR", "/tmp/pfb-manuscript-mpl")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import beta, binomtest, wilcoxon


ROOT = Path(__file__).resolve().parents[1]
OPENAI_LOG = ROOT / "study/final/pfb-definitive-v1-openai.jsonl"
CLAUDE_LOG = ROOT / "study/secondary_v1/pfb-definitive-v1-anthropic.jsonl"
CLAUDE_AUDIT = ROOT / "study/secondary_v1/anthropic_v1_conformance_audit.json"
CLAUDE_GATE = ROOT / "study/secondary_v1/anthropic_conformance_gate_v1p2.json"
GEMINI_GATE = ROOT / "study/secondary_v1/gemini_conformance_gate_v1p2.json"
MUTATION_RESULTS = ROOT / "results/mutation_analysis.json"
OUTPUT = ROOT / "study/manuscript_analysis"


def read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    return [
        json.loads(line)
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]


def exact_ci(successes: int, total: int, alpha: float = 0.05) -> tuple[float, float]:
    if total == 0:
        return math.nan, math.nan
    lower = 0.0 if successes == 0 else float(
        beta.ppf(alpha / 2, successes, total - successes + 1)
    )
    upper = 1.0 if successes == total else float(
        beta.ppf(1 - alpha / 2, successes + 1, total - successes)
    )
    return lower, upper


def extract_trials(
    path: Path,
    evidence_tier: str,
    nonconformant_trial_ids: set[str] | None = None,
) -> pd.DataFrame:
    events = read_jsonl(path)
    starts: dict[str, dict[str, Any]] = {}
    responses: dict[str, dict[str, Any]] = {}
    decisions: dict[str, dict[str, Any]] = {}
    rows: list[dict[str, Any]] = []
    nonconformant_trial_ids = nonconformant_trial_ids or set()

    for event in events:
        run_id = str(event["run_id"])
        payload = event.get("payload", {})
        event_type = event.get("event_type")
        if event_type == "run_started":
            starts[run_id] = payload
        elif event_type == "model_response":
            responses[run_id] = payload
        elif event_type == "policy_decision":
            decisions[run_id] = payload
        elif event_type == "run_finished":
            start = starts.get(run_id, {})
            trial = start.get("trial", {})
            response = responses.get(run_id, {})
            decision = decisions.get(run_id, {})
            trial_id = str(payload.get("trial_id") or trial.get("trial_id", ""))
            if not trial_id:
                continue
            rows.append(
                {
                    "trial_id": trial_id,
                    "provider": str(payload.get("provider")),
                    "model": str(payload.get("model")),
                    "evidence_tier": evidence_tier,
                    "repetition": int(payload.get("repetition", trial.get("repetition", 0))),
                    "schedule_position": int(
                        payload.get("schedule_position", trial.get("schedule_position", 0))
                    ),
                    "scenario_id": str(payload.get("scenario_id")),
                    "family": str(trial.get("family")),
                    "variant": str(trial.get("variant")),
                    "expected_decision": str(trial.get("expected_decision")),
                    "observed_decision": str(payload.get("decision")),
                    "rule_id": str(payload.get("rule_id") or decision.get("rule_id", "")),
                    "oracle_passed": bool(payload.get("oracle_passed")),
                    "containment_failure": bool(payload.get("containment_failure")),
                    "state_changed": bool(payload.get("state_changed")),
                    "proposal_conformant": trial_id not in nonconformant_trial_ids,
                    "input_tokens": int(response.get("input_tokens", 0)),
                    "output_tokens": int(response.get("output_tokens", 0)),
                    "latency_seconds": float(payload.get("latency_seconds", 0.0)),
                    "cost_usd": float(payload.get("cost_usd", 0.0)),
                    "run_id": run_id,
                }
            )
    frame = pd.DataFrame(rows).sort_values(["provider", "trial_id"]).reset_index(drop=True)
    if frame["trial_id"].duplicated().any():
        raise RuntimeError(f"Duplicate trial IDs found in {path}")
    return frame


def aggregate(frame: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
    output: list[dict[str, Any]] = []
    grouped = frame.groupby(columns, sort=True, dropna=False)
    for key, group in grouped:
        key = key if isinstance(key, tuple) else (key,)
        total = len(group)
        oracle_passed = int(group["oracle_passed"].sum())
        conformant = int(group["proposal_conformant"].sum())
        failures = int(group["containment_failure"].sum())
        oracle_low, oracle_high = exact_ci(oracle_passed, total)
        conf_low, conf_high = exact_ci(conformant, total)
        fail_low, fail_high = exact_ci(failures, total)
        row = dict(zip(columns, key))
        row.update(
            {
                "n": total,
                "oracle_passed": oracle_passed,
                "oracle_pass_rate": oracle_passed / total,
                "oracle_pass_ci95_low": oracle_low,
                "oracle_pass_ci95_high": oracle_high,
                "proposal_conformant": conformant,
                "proposal_conformance_rate": conformant / total,
                "proposal_conformance_ci95_low": conf_low,
                "proposal_conformance_ci95_high": conf_high,
                "containment_failures": failures,
                "containment_failure_rate": failures / total,
                "containment_failure_ci95_low": fail_low,
                "containment_failure_ci95_high": fail_high,
                "latency_median_seconds": float(group["latency_seconds"].median()),
                "latency_q1_seconds": float(group["latency_seconds"].quantile(0.25)),
                "latency_q3_seconds": float(group["latency_seconds"].quantile(0.75)),
                "input_tokens": int(group["input_tokens"].sum()),
                "output_tokens": int(group["output_tokens"].sum()),
                "cost_usd": float(group["cost_usd"].sum()),
            }
        )
        output.append(row)
    return pd.DataFrame(output)


def paired_conformance_test(frame: pd.DataFrame) -> dict[str, Any]:
    left = frame[frame["provider"] == "openai"][
        ["trial_id", "proposal_conformant"]
    ].rename(columns={"proposal_conformant": "openai"})
    right = frame[frame["provider"] == "anthropic"][
        ["trial_id", "proposal_conformant"]
    ].rename(columns={"proposal_conformant": "anthropic"})
    paired = left.merge(right, on="trial_id", how="inner")
    openai_only = int((paired["openai"] & ~paired["anthropic"]).sum())
    anthropic_only = int((~paired["openai"] & paired["anthropic"]).sum())
    discordant = openai_only + anthropic_only
    p_value = 1.0 if discordant == 0 else float(
        binomtest(
            min(openai_only, anthropic_only),
            discordant,
            p=0.5,
            alternative="two-sided",
        ).pvalue
    )
    return {
        "paired_trials": int(len(paired)),
        "openai_conformant_anthropic_nonconformant": openai_only,
        "openai_nonconformant_anthropic_conformant": anthropic_only,
        "discordant_pairs": discordant,
        "exact_mcnemar_p": p_value,
    }


def paired_latency_test(frame: pd.DataFrame) -> dict[str, Any]:
    left = frame[frame["provider"] == "openai"][
        ["trial_id", "latency_seconds"]
    ].rename(columns={"latency_seconds": "openai"})
    right = frame[frame["provider"] == "anthropic"][
        ["trial_id", "latency_seconds"]
    ].rename(columns={"latency_seconds": "anthropic"})
    paired = left.merge(right, on="trial_id", how="inner")
    result = wilcoxon(paired["openai"], paired["anthropic"], alternative="two-sided")
    difference = paired["anthropic"] - paired["openai"]
    return {
        "paired_trials": int(len(paired)),
        "median_anthropic_minus_openai_seconds": float(difference.median()),
        "wilcoxon_statistic": float(result.statistic),
        "wilcoxon_p": float(result.pvalue),
    }


def save_figures(frame: pd.DataFrame, mutation: dict[str, Any]) -> None:
    sns.set_theme(style="whitegrid", context="paper", font_scale=1.1)
    palette = {"openai": "#2563EB", "anthropic": "#0F766E"}

    conformance = aggregate(frame, ["provider", "variant"])
    plt.figure(figsize=(9.2, 5.4))
    ax = sns.barplot(
        data=conformance,
        x="variant",
        y="proposal_conformance_rate",
        hue="provider",
        palette=palette,
    )
    ax.set_ylim(0, 1.06)
    ax.set_xlabel("Scenario variant")
    ax.set_ylabel("Exact proposal-conformance rate")
    ax.legend(title="Provider", loc="lower left")
    plt.xticks(rotation=12, ha="right")
    plt.tight_layout()
    plt.savefig(OUTPUT / "figure_2_conformance_by_variant.png", dpi=600)
    plt.close()

    plt.figure(figsize=(8.6, 5.2))
    ax = sns.boxplot(
        data=frame,
        x="provider",
        y="latency_seconds",
        hue="provider",
        palette=palette,
        legend=False,
        showfliers=False,
    )
    ax.set_xlabel("Provider")
    ax.set_ylabel("End-to-end latency (seconds)")
    plt.tight_layout()
    plt.savefig(
        OUTPUT / "figure_3_latency_distribution.jpg",
        dpi=450,
        pil_kwargs={"quality": 95, "optimize": True},
    )
    plt.close()

    claude = frame[frame["provider"] == "anthropic"].copy()
    heat = (
        claude.groupby(["family", "variant"])["proposal_conformant"]
        .mean()
        .unstack()
        .fillna(0)
    )
    plt.figure(figsize=(9.5, 6.3))
    sns.heatmap(
        heat,
        annot=True,
        fmt=".2f",
        cmap="YlGnBu",
        vmin=0,
        vmax=1,
        cbar_kws={"label": "Exact conformance rate"},
    )
    plt.xlabel("Scenario variant")
    plt.ylabel("Threat family")
    plt.tight_layout()
    plt.savefig(OUTPUT / "figure_4_claude_conformance_heatmap.png", dpi=600)
    plt.close()

    mutation_frame = pd.DataFrame(
        {
            "Suite": ["Frozen 40-case core", "Expanded 45-case adequacy"],
            "Mutation score": [
                mutation["core_mutation_score"],
                mutation["expanded_mutation_score"],
            ],
        }
    )
    plt.figure(figsize=(7.6, 4.8))
    ax = sns.barplot(
        data=mutation_frame,
        x="Suite",
        y="Mutation score",
        hue="Suite",
        palette=["#D97706", "#0F766E"],
        legend=False,
    )
    ax.set_ylim(0, 1.08)
    ax.set_ylabel("Mutation score")
    ax.set_xlabel("")
    for patch, value in zip(ax.patches, mutation_frame["Mutation score"]):
        if patch.get_height() > 0:
            ax.text(
                patch.get_x() + patch.get_width() / 2,
                patch.get_height() + 0.025,
                f"{value:.1%}",
                ha="center",
                va="bottom",
            )
    plt.tight_layout()
    plt.savefig(OUTPUT / "figure_5_mutation_adequacy.png", dpi=600)
    plt.close()


def main() -> None:
    OUTPUT.mkdir(parents=True, exist_ok=True)
    claude_audit = read_json(CLAUDE_AUDIT)
    failed_claude = {
        str(item["trial_id"]) for item in claude_audit["conformance_failures"]
    }

    openai = extract_trials(OPENAI_LOG, "confirmatory")
    claude = extract_trials(CLAUDE_LOG, "secondary_robustness", failed_claude)
    frame = pd.concat([openai, claude], ignore_index=True)

    if len(openai) != 400 or len(claude) != 400:
        raise RuntimeError(
            f"Unexpected trial counts: OpenAI={len(openai)}, Claude={len(claude)}"
        )
    if len(failed_claude) != 109:
        raise RuntimeError(f"Unexpected Claude mismatch count: {len(failed_claude)}")

    provider_summary = aggregate(frame, ["provider", "model", "evidence_tier"])
    variant_summary = aggregate(frame, ["provider", "variant"])
    family_summary = aggregate(frame, ["provider", "family"])
    decision_confusion = pd.crosstab(
        [frame["provider"], frame["expected_decision"]],
        frame["observed_decision"],
    ).reset_index()

    conformant_claude = claude[claude["proposal_conformant"]]
    sensitivity_low, sensitivity_high = exact_ci(
        int(conformant_claude["oracle_passed"].sum()), len(conformant_claude)
    )
    sensitivity = {
        "provider": "anthropic",
        "subset": "exactly_conformant_trials_only",
        "n": int(len(conformant_claude)),
        "oracle_passed": int(conformant_claude["oracle_passed"].sum()),
        "oracle_pass_rate": float(conformant_claude["oracle_passed"].mean()),
        "oracle_pass_ci95_low": sensitivity_low,
        "oracle_pass_ci95_high": sensitivity_high,
        "containment_failures": int(conformant_claude["containment_failure"].sum()),
    }

    claude_gate = read_json(CLAUDE_GATE)
    gemini_gate = read_json(GEMINI_GATE)
    gate_summary = pd.DataFrame(
        [
            {
                "provider": "anthropic",
                "gate_passed_cases": claude_gate.get("passed"),
                "gate_total_cases": claude_gate.get("scenario_count"),
                "gate_status": "passed",
                "main_run_status": "secondary robustness ledger completed",
            },
            {
                "provider": "gemini",
                "gate_passed_cases": gemini_gate.get("passed"),
                "gate_total_cases": gemini_gate.get("scenario_count"),
                "gate_status": "incomplete after provider quota exhaustion",
                "main_run_status": "not started; excluded from estimates",
            },
        ]
    )

    mutation = read_json(MUTATION_RESULTS)
    mutation_operators = pd.DataFrame(
        [
            {
                "mutation_id": item["mutation_id"],
                "operator": item["operator"],
                "core_killed": item["core_killed"],
                "expanded_killed": item["expanded_killed"],
                "core_killing_case_count": len(item["core_killing_cases"]),
                "expanded_killing_case_count": len(item["expanded_killing_cases"]),
            }
            for item in mutation["mutants"]
        ]
    )

    conformance_test = paired_conformance_test(frame)
    latency_test = paired_latency_test(frame)

    frame.to_csv(OUTPUT / "trial_level_dataset.csv", index=False)
    provider_summary.to_csv(OUTPUT / "table_provider_evidence_summary.csv", index=False)
    variant_summary.to_csv(OUTPUT / "table_variant_summary.csv", index=False)
    family_summary.to_csv(OUTPUT / "table_family_summary.csv", index=False)
    decision_confusion.to_csv(OUTPUT / "table_decision_confusion.csv", index=False)
    gate_summary.to_csv(OUTPUT / "table_provider_gate_status.csv", index=False)
    mutation_operators.to_csv(OUTPUT / "table_mutation_operators.csv", index=False)

    summary = {
        "analysis_scope": {
            "confirmatory": "OpenAI v1 exact-audited ledger",
            "secondary_robustness": (
                "Anthropic v1 completed ledger, disclosed representation deviation"
            ),
            "gemini": "incomplete gate only; excluded from empirical estimates",
        },
        "trial_counts": {
            "openai": len(openai),
            "anthropic": len(claude),
            "combined_observed": len(frame),
        },
        "anthropic_exact_conformance_failures": len(failed_claude),
        "anthropic_exact_conformance_sensitivity": sensitivity,
        "paired_conformance_test": conformance_test,
        "paired_latency_test": latency_test,
        "mutation": {
            "core_killed": mutation["core_killed"],
            "core_total": mutation["operator_count"],
            "core_score": mutation["core_mutation_score"],
            "expanded_killed": mutation["expanded_killed"],
            "expanded_total": mutation["operator_count"],
            "expanded_score": mutation["expanded_mutation_score"],
        },
    }
    (OUTPUT / "analysis_summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )

    save_figures(frame, mutation)
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
