from __future__ import annotations

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

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

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


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 clopper_pearson(successes: int, total: int, alpha: float = 0.05):
    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_provider(path: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    events = read_jsonl(path)
    starts: dict[str, dict[str, Any]] = {}
    responses: dict[str, dict[str, Any]] = {}
    decisions: dict[str, dict[str, Any]] = {}
    rows: dict[str, dict[str, Any]] = {}
    diagnostics = {
        "event_log": path.name,
        "transient_aborts": 0,
        "harness_aborts": 0,
        "terminal_model_conformance_failures": 0,
    }
    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":
            trial = starts.get(run_id, {}).get("trial", {})
            trial_id = str(payload.get("trial_id") or trial.get("trial_id", ""))
            if not trial_id:
                continue
            response = responses.get(run_id, {})
            decision = decisions.get(run_id, {})
            rows[trial_id] = {
                "trial_id": trial_id,
                "provider": payload.get("provider"),
                "model": payload.get("model"),
                "repetition": int(
                    payload.get("repetition") or trial.get("repetition", 0)
                ),
                "schedule_position": int(
                    payload.get("schedule_position")
                    or trial.get("schedule_position", 0)
                ),
                "scenario_id": payload.get("scenario_id"),
                "family": trial.get("family"),
                "variant": trial.get("variant"),
                "expected_decision": trial.get("expected_decision"),
                "observed_decision": payload.get("decision"),
                "rule_id": 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": True,
                "terminal_abort": False,
                "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,
            }
        elif event_type == "run_aborted":
            abort_class = str(payload.get("abort_class", "harness_error"))
            if abort_class == "provider_transient":
                diagnostics["transient_aborts"] += 1
            elif abort_class == "model_conformance":
                diagnostics["terminal_model_conformance_failures"] += 1
                trial = payload.get("trial", {})
                trial_id = str(trial.get("trial_id", ""))
                if trial_id:
                    rows[trial_id] = {
                        "trial_id": trial_id,
                        "provider": payload.get("provider"),
                        "model": payload.get("model"),
                        "repetition": int(trial.get("repetition", 0)),
                        "schedule_position": int(
                            trial.get("schedule_position", 0)
                        ),
                        "scenario_id": payload.get("scenario_id"),
                        "family": trial.get("family"),
                        "variant": trial.get("variant"),
                        "expected_decision": trial.get("expected_decision"),
                        "observed_decision": "conformance_abort",
                        "rule_id": "",
                        "oracle_passed": False,
                        "containment_failure": False,
                        "state_changed": False,
                        "proposal_conformant": False,
                        "terminal_abort": True,
                        "input_tokens": 0,
                        "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,
                    }
            else:
                diagnostics["harness_aborts"] += 1
    return list(rows.values()), diagnostics


def aggregate_group(frame: pd.DataFrame, group_columns: list[str]) -> pd.DataFrame:
    output = []
    grouped = frame.groupby(group_columns, dropna=False, sort=True)
    for key, group in grouped:
        key = key if isinstance(key, tuple) else (key,)
        total = len(group)
        passed = int(group["oracle_passed"].sum())
        conformant = int(group["proposal_conformant"].sum())
        failures = int(group["containment_failure"].sum())
        pass_low, pass_high = clopper_pearson(passed, total)
        fail_low, fail_high = clopper_pearson(failures, total)
        row = dict(zip(group_columns, key))
        row.update(
            {
                "n": total,
                "oracle_passed": passed,
                "oracle_pass_rate": passed / total if total else math.nan,
                "oracle_pass_ci95_low": pass_low,
                "oracle_pass_ci95_high": pass_high,
                "proposal_conformance_rate": (
                    conformant / total if total else math.nan
                ),
                "containment_failures": failures,
                "containment_failure_rate": (
                    failures / total if total else math.nan
                ),
                "containment_failure_ci95_low": fail_low,
                "containment_failure_ci95_high": fail_high,
                "latency_median_seconds": float(
                    group["latency_seconds"].median()
                ),
                "latency_iqr_seconds": float(
                    group["latency_seconds"].quantile(0.75)
                    - group["latency_seconds"].quantile(0.25)
                ),
                "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 holm_adjust(p_values: list[float]) -> list[float]:
    indexed = sorted(enumerate(p_values), key=lambda item: item[1])
    adjusted = [1.0] * len(p_values)
    running = 0.0
    count = len(p_values)
    for rank, (index, value) in enumerate(indexed):
        candidate = min(1.0, (count - rank) * value)
        running = max(running, candidate)
        adjusted[index] = running
    return adjusted


def paired_model_tests(frame: pd.DataFrame) -> pd.DataFrame:
    providers = sorted(frame["provider"].dropna().unique())
    results = []
    raw_p = []
    for left_index, left in enumerate(providers):
        for right in providers[left_index + 1 :]:
            a = frame[frame["provider"] == left][
                ["trial_id", "oracle_passed"]
            ].rename(columns={"oracle_passed": "left_pass"})
            b = frame[frame["provider"] == right][
                ["trial_id", "oracle_passed"]
            ].rename(columns={"oracle_passed": "right_pass"})
            paired = a.merge(b, on="trial_id", how="inner")
            left_only = int(
                ((paired["left_pass"]) & (~paired["right_pass"])).sum()
            )
            right_only = int(
                ((~paired["left_pass"]) & (paired["right_pass"])).sum()
            )
            discordant = left_only + right_only
            p_value = (
                1.0
                if discordant == 0
                else float(
                    binomtest(
                        min(left_only, right_only),
                        discordant,
                        0.5,
                        alternative="two-sided",
                    ).pvalue
                )
            )
            results.append(
                {
                    "provider_a": left,
                    "provider_b": right,
                    "paired_trials": len(paired),
                    "a_pass_b_fail": left_only,
                    "a_fail_b_pass": right_only,
                    "mcnemar_exact_p": p_value,
                }
            )
            raw_p.append(p_value)
    adjusted = holm_adjust(raw_p)
    for row, value in zip(results, adjusted):
        row["holm_adjusted_p"] = value
    return pd.DataFrame(results)


def latency_test(frame: pd.DataFrame) -> dict[str, Any]:
    groups = [
        group["latency_seconds"].to_numpy()
        for _, group in frame.groupby("provider")
        if len(group)
    ]
    if len(groups) < 2:
        return {"test": "Kruskal-Wallis", "statistic": None, "p_value": None}
    result = kruskal(*groups)
    return {
        "test": "Kruskal-Wallis",
        "statistic": float(result.statistic),
        "p_value": float(result.pvalue),
    }


def fit_gee_if_estimable(frame: pd.DataFrame, output_dir: Path) -> dict[str, Any]:
    if frame["oracle_passed"].nunique() < 2:
        result = {
            "status": "not_estimable",
            "reason": "The binary endpoint has no variation.",
        }
    else:
        try:
            import statsmodels.api as sm
            import statsmodels.formula.api as smf

            model = smf.gee(
                "oracle_passed ~ C(provider) + C(variant)",
                groups="family",
                data=frame.assign(
                    oracle_passed=frame["oracle_passed"].astype(int)
                ),
                family=sm.families.Binomial(),
            )
            fitted = model.fit()
            (output_dir / "gee_logistic_summary.txt").write_text(
                fitted.summary().as_text() + "\n", encoding="utf-8"
            )
            result = {
                "status": "fitted",
                "formula": "oracle_passed ~ C(provider) + C(variant)",
                "cluster": "family",
            }
        except Exception as exc:
            result = {
                "status": "failed",
                "reason": f"{type(exc).__name__}: {exc}",
            }
    (output_dir / "gee_logistic_status.json").write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    return result


def save_figures(frame: pd.DataFrame, output_dir: Path) -> None:
    sns.set_theme(style="whitegrid", context="paper")

    by_variant = aggregate_group(frame, ["provider", "variant"])
    plt.figure(figsize=(9, 5.5))
    sns.barplot(
        data=by_variant,
        x="variant",
        y="oracle_pass_rate",
        hue="provider",
    )
    plt.ylim(0, 1.05)
    plt.ylabel("Oracle pass rate")
    plt.xlabel("Scenario variant")
    plt.tight_layout()
    plt.savefig(output_dir / "figure_pass_rate_by_variant.png", dpi=300)
    plt.close()

    plt.figure(figsize=(8, 5))
    sns.boxplot(data=frame, x="provider", y="latency_seconds")
    plt.ylabel("Latency (seconds)")
    plt.xlabel("Provider")
    plt.tight_layout()
    plt.savefig(output_dir / "figure_latency_by_provider.png", dpi=300)
    plt.close()

    heat = (
        frame.groupby(["family", "provider"])["oracle_passed"]
        .mean()
        .unstack()
    )
    plt.figure(figsize=(8, 7))
    sns.heatmap(
        heat,
        vmin=0,
        vmax=1,
        annot=True,
        fmt=".2f",
        cmap="Blues",
        cbar_kws={"label": "Oracle pass rate"},
    )
    plt.xlabel("Provider")
    plt.ylabel("Threat family")
    plt.tight_layout()
    plt.savefig(output_dir / "figure_family_heatmap.png", dpi=300)
    plt.close()

    cost = frame.groupby("provider", as_index=False)["cost_usd"].sum()
    plt.figure(figsize=(7, 4.5))
    sns.barplot(data=cost, x="provider", y="cost_usd")
    plt.ylabel("Logged cost (USD)")
    plt.xlabel("Provider")
    plt.tight_layout()
    plt.savefig(output_dir / "figure_cost_by_provider.png", dpi=300)
    plt.close()


def dataframe_markdown(frame: pd.DataFrame) -> str:
    columns = list(frame.columns)
    lines = [
        "| " + " | ".join(columns) + " |",
        "|" + "|".join("---" for _ in columns) + "|",
    ]
    for _, row in frame.iterrows():
        values = []
        for column in columns:
            value = row[column]
            if isinstance(value, float):
                values.append(f"{value:.6f}")
            else:
                values.append(str(value))
        lines.append("| " + " | ".join(values) + " |")
    return "\n".join(lines)


def write_markdown_report(
    frame: pd.DataFrame,
    provider_summary: pd.DataFrame,
    diagnostics: list[dict[str, Any]],
    latency: dict[str, Any],
    gee: dict[str, Any],
    output_dir: Path,
    planned_per_provider: int,
) -> None:
    providers = sorted(frame["provider"].dropna().unique())
    complete = all(
        len(frame[frame["provider"] == provider]) == planned_per_provider
        for provider in providers
    )
    lines = [
        "# PolicyFaultBench Definitive Study — Automated Analysis",
        "",
        f"**Analysis status:** {'complete' if complete else 'partial'}",
        "",
        "## Provider-level primary endpoints",
        "",
        dataframe_markdown(provider_summary),
        "",
        "## Diagnostics",
        "",
        "```json",
        json.dumps(diagnostics, indent=2, sort_keys=True),
        "```",
        "",
        "## Inferential analyses",
        "",
        f"- Latency comparison: `{json.dumps(latency, sort_keys=True)}`",
        f"- Clustered binary model: `{json.dumps(gee, sort_keys=True)}`",
        "",
        "Primary model comparisons use paired exact McNemar tests with Holm "
        "correction. Binary confidence intervals are exact 95% "
        "Clopper–Pearson intervals.",
        "",
        "Harness/provider transient retries are diagnostic events and are not "
        "silently converted into successful primary outcomes. Terminal model "
        "proposal-conformance failures remain primary failures.",
        "",
    ]
    (output_dir / "FINAL_STUDY_ANALYSIS.md").write_text(
        "\n".join(lines), encoding="utf-8"
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--study-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--planned-per-provider", type=int, default=400)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    event_logs = sorted(args.study_dir.glob("pfb-definitive-v1-*.jsonl"))
    if not event_logs:
        raise SystemExit("No definitive-study JSONL files were found")
    all_rows = []
    diagnostics = []
    for path in event_logs:
        rows, diagnostic = extract_provider(path)
        all_rows.extend(rows)
        diagnostics.append(diagnostic)
    frame = pd.DataFrame(all_rows).sort_values(
        ["provider", "schedule_position"]
    )
    frame.to_csv(args.output_dir / "trial_results.csv", index=False)

    provider_summary = aggregate_group(frame, ["provider"])
    variant_summary = aggregate_group(frame, ["provider", "variant"])
    family_summary = aggregate_group(frame, ["provider", "family"])
    provider_summary.to_csv(
        args.output_dir / "provider_summary.csv", index=False
    )
    variant_summary.to_csv(
        args.output_dir / "provider_variant_summary.csv", index=False
    )
    family_summary.to_csv(
        args.output_dir / "provider_family_summary.csv", index=False
    )

    confusion = (
        frame.groupby(
            ["provider", "expected_decision", "observed_decision"],
            dropna=False,
        )
        .size()
        .reset_index(name="count")
    )
    confusion.to_csv(args.output_dir / "decision_confusion_matrix.csv", index=False)
    paired = paired_model_tests(frame)
    paired.to_csv(args.output_dir / "paired_model_tests.csv", index=False)
    latency = latency_test(frame)
    gee = fit_gee_if_estimable(frame, args.output_dir)
    save_figures(frame, args.output_dir)

    machine_summary = {
        "schema_version": "1.0",
        "providers_found": sorted(frame["provider"].dropna().unique()),
        "primary_rows": len(frame),
        "planned_per_provider": args.planned_per_provider,
        "provider_counts": {
            str(key): int(value)
            for key, value in frame.groupby("provider").size().items()
        },
        "diagnostics": diagnostics,
        "latency_test": latency,
        "gee": gee,
    }
    (args.output_dir / "analysis_summary.json").write_text(
        json.dumps(machine_summary, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    write_markdown_report(
        frame,
        provider_summary,
        diagnostics,
        latency,
        gee,
        args.output_dir,
        args.planned_per_provider,
    )
    print(json.dumps(machine_summary, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
