from __future__ import annotations

import json
import time
from dataclasses import dataclass
from itertools import pairwise
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt

from lost_cities_jax.ppo import (
    evaluate_checkpoint_match,
    evaluate_checkpoint_vs_static,
    load_config,
)

DATE = "2026-07-05"
ARTIFACT_DIR = Path("/mnt/2tbhdd/coolrl-lost-cities-artifacts/diminishing-returns/2026-07-05")
REPORT_PATH = Path(f"docs/reports/diminishing-returns-{DATE}.md")
SUMMARY_PATH = Path(f"docs/reports/diminishing-returns-{DATE}-summary.jsonl")
GAMES_PER_SEAT = 2000
TOTAL_DUPLICATE_GAMES = GAMES_PER_SEAT * 2


@dataclass(frozen=True)
class CheckpointSpec:
    name: str
    label: str
    config: str
    checkpoint: str
    stage_index: int


CHECKPOINTS = [
    CheckpointSpec(
        name="ladder_v2_gate3",
        label="ladder v2 gate3",
        config="configs/jax_ppo/ladder-v2-expert.yaml",
        checkpoint="/mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_013223_jax-ppo-ladder-v2-expert/latest",
        stage_index=0,
    ),
    CheckpointSpec(
        name="league_v1_update_250",
        label="league v1 update 250",
        config="/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/main_ppo_config.json",
        checkpoint="/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250",
        stage_index=1,
    ),
    CheckpointSpec(
        name="league_v1_update_500",
        label="league v1 update 500",
        config="/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/main_ppo_config.json",
        checkpoint="/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500",
        stage_index=2,
    ),
    CheckpointSpec(
        name="repair_c01_update_500",
        label="repair c01 update 500",
        config="/mnt/2tbhdd/coolrl-lost-cities-artifacts/gates-1-2/2026-07-05_094238_gates-1-2/gate2c/league/2026-07-05_125709_jax-ppo-gates-1-2-repair-c01/main_ppo_config.json",
        checkpoint="/mnt/2tbhdd/coolrl-lost-cities-artifacts/gates-1-2/2026-07-05_094238_gates-1-2/gate2c/league/2026-07-05_125709_jax-ppo-gates-1-2-repair-c01/snapshots/cycle_01_update_000500",
        stage_index=3,
    ),
]


EXPLOITER_ROWS = [
    {
        "phase": "ladder_v2",
        "target": "ladder_v2_gate3",
        "protocol": "short_random_shaping",
        "budget": "same-budget phase-1 exploiter",
        "warmstart": "random",
        "win_rate": 0.83135,
        "mean_score_diff": 54.7721,
        "source": "docs/reports/ladder-v2-2026-07-05-summary.jsonl",
    },
    {
        "phase": "league_v1",
        "target": "league_v1_update_250",
        "protocol": "league_cycle_exploiter",
        "budget": "league v1 cycle exploiter",
        "warmstart": "random",
        "win_rate": 0.50225,
        "mean_score_diff": 0.6818,
        "source": "docs/reports/league-v1-2026-07-05-summary.jsonl",
    },
    {
        "phase": "league_v1",
        "target": "league_v1_update_500",
        "protocol": "league_cycle_exploiter",
        "budget": "same exploiter vs final",
        "warmstart": "random",
        "win_rate": 0.47095,
        "mean_score_diff": -2.63775,
        "source": "docs/reports/league-v1-2026-07-05-summary.jsonl",
    },
    {
        "phase": "gates_1_2_original",
        "target": "league_v1_update_500",
        "protocol": "long_random_shaping",
        "budget": "1200 updates, shaping anneal",
        "warmstart": "random",
        "win_rate": 0.587,
        "mean_score_diff": 10.95125,
        "source": "docs/reports/gates-1-2-2026-07-05-summary.jsonl",
    },
    {
        "phase": "gates_1_2_original",
        "target": "league_v1_update_500",
        "protocol": "warmstart_gate3_no_shaping",
        "budget": "900 updates, shaping 0",
        "warmstart": "ladder_v2_gate3",
        "win_rate": 0.590,
        "mean_score_diff": 11.38075,
        "source": "docs/reports/gates-1-2-2026-07-05-summary.jsonl",
    },
    {
        "phase": "gates_1_2_original",
        "target": "league_v1_update_500",
        "protocol": "replay_exploiter_no_shaping",
        "budget": "900 updates, shaping 0",
        "warmstart": "league_v1_cycle_1_exploiter",
        "win_rate": 0.582,
        "mean_score_diff": 10.42525,
        "source": "docs/reports/gates-1-2-2026-07-05-summary.jsonl",
    },
    {
        "phase": "repair_c01",
        "target": "repair_c01_update_500",
        "protocol": "long_random_shaping",
        "budget": "1200 updates, shaping anneal",
        "warmstart": "random",
        "win_rate": 0.539,
        "mean_score_diff": 6.27925,
        "source": "docs/reports/gates-1-2-2026-07-05-summary.jsonl",
    },
    {
        "phase": "repair_c01",
        "target": "repair_c01_update_500",
        "protocol": "warmstart_gate3_no_shaping",
        "budget": "900 updates, shaping 0",
        "warmstart": "ladder_v2_gate3",
        "win_rate": 0.549,
        "mean_score_diff": 7.0285,
        "source": "docs/reports/gates-1-2-2026-07-05-summary.jsonl",
    },
    {
        "phase": "repair_c01",
        "target": "repair_c01_update_500",
        "protocol": "replay_exploiter_no_shaping",
        "budget": "900 updates, shaping 0",
        "warmstart": "league_v1_cycle_1_exploiter",
        "win_rate": 0.53925,
        "mean_score_diff": 6.10225,
        "source": "docs/reports/gates-1-2-2026-07-05-summary.jsonl",
    },
]


def main() -> None:
    started = time.perf_counter()
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)

    h2h_rows = run_h2h()
    expert_rows = run_expert_anchor()
    exploiter_rows = add_exploiter_deltas(EXPLOITER_ROWS)
    plot_paths = write_plots(h2h_rows, expert_rows, exploiter_rows)
    report = build_report(h2h_rows, expert_rows, exploiter_rows, plot_paths)
    REPORT_PATH.write_text(report, encoding="utf-8")
    write_summary(h2h_rows, expert_rows, exploiter_rows, time.perf_counter() - started)


def run_h2h() -> list[dict[str, Any]]:
    rows = []
    for previous, current in pairwise(CHECKPOINTS):
        output = ARTIFACT_DIR / f"h2h_{current.name}_vs_{previous.name}.json"
        result = evaluate_checkpoint_match(
            load_config(current.config),
            current.checkpoint,
            load_config(previous.config),
            previous.checkpoint,
            games=GAMES_PER_SEAT,
            duplicate=True,
            output=output,
        )
        rows.append(
            {
                "event": "h2h_adjacent",
                "previous": previous.name,
                "current": current.name,
                "previous_label": previous.label,
                "current_label": current.label,
                "json": str(output),
                **summarize_match(result),
            }
        )
    return rows


def run_expert_anchor() -> list[dict[str, Any]]:
    rows = []
    for spec in CHECKPOINTS:
        output = ARTIFACT_DIR / f"expert_anchor_{spec.name}.json"
        result = evaluate_checkpoint_vs_static(
            load_config(spec.config),
            spec.checkpoint,
            "heuristic_expert",
            games=GAMES_PER_SEAT,
            duplicate=True,
            output=output,
        )
        rows.append(
            {
                "event": "expert_anchor",
                "checkpoint": spec.name,
                "label": spec.label,
                "stage_index": spec.stage_index,
                "json": str(output),
                **summarize_match(result),
            }
        )
    return rows


def summarize_match(result: dict[str, Any]) -> dict[str, Any]:
    return {
        "games": result["games"],
        "win_rate": result["win_rate"],
        "wilson_low": result["wilson_low"],
        "wilson_high": result["wilson_high"],
        "mean_score_diff": result["mean_score_diff"],
        "score_diff_ci95_low": result["score_diff_ci95_low"],
        "score_diff_ci95_high": result["score_diff_ci95_high"],
        "score_diff_std": result["score_diff_std"],
        "mean_game_length": result["mean_game_length"],
        "max_steps_rate": result["max_steps_rate"],
        "opened_colors_per_game": result["opened_colors_per_game"],
        "play_action_rate": result["play_action_rate"],
        "positive_expeditions_per_game": result["positive_expeditions_per_game"],
    }


def add_exploiter_deltas(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    enriched = [dict(row) for row in rows]
    by_protocol: dict[str, list[dict[str, Any]]] = {}
    for row in enriched:
        by_protocol.setdefault(row["protocol"], []).append(row)
    for protocol_rows in by_protocol.values():
        previous = None
        for row in protocol_rows:
            row["delta_from_previous_same_protocol"] = (
                None if previous is None else row["win_rate"] - previous["win_rate"]
            )
            previous = row
    return enriched


def write_plots(
    h2h_rows: list[dict[str, Any]],
    expert_rows: list[dict[str, Any]],
    exploiter_rows: list[dict[str, Any]],
) -> dict[str, str]:
    paths = {
        "h2h": REPORT_PATH.parent / f"diminishing-returns-{DATE}-h2h.png",
        "expert": REPORT_PATH.parent / f"diminishing-returns-{DATE}-expert.png",
        "exploiter": REPORT_PATH.parent / f"diminishing-returns-{DATE}-exploiter.png",
    }

    plt.figure(figsize=(7.5, 4.2))
    x = list(range(len(h2h_rows)))
    means = [row["mean_score_diff"] for row in h2h_rows]
    lows = [row["score_diff_ci95_low"] for row in h2h_rows]
    highs = [row["score_diff_ci95_high"] for row in h2h_rows]
    plt.errorbar(
        x,
        means,
        yerr=[
            [m - lo for m, lo in zip(means, lows, strict=True)],
            [hi - m for m, hi in zip(means, highs, strict=True)],
        ],
        fmt="o-",
        capsize=4,
    )
    plt.axhline(0, color="black", linewidth=1)
    plt.xticks(x, [row["current"] for row in h2h_rows], rotation=20, ha="right")
    plt.ylabel("Mean score diff vs previous")
    plt.title("Adjacent checkpoint gains")
    plt.tight_layout()
    plt.savefig(paths["h2h"], dpi=160)
    plt.close()

    plt.figure(figsize=(7.5, 4.2))
    x = [row["stage_index"] for row in expert_rows]
    means = [row["mean_score_diff"] for row in expert_rows]
    lows = [row["score_diff_ci95_low"] for row in expert_rows]
    highs = [row["score_diff_ci95_high"] for row in expert_rows]
    plt.errorbar(
        x,
        means,
        yerr=[
            [m - lo for m, lo in zip(means, lows, strict=True)],
            [hi - m for m, hi in zip(means, highs, strict=True)],
        ],
        fmt="o-",
        capsize=4,
    )
    plt.axhline(0, color="black", linewidth=1)
    plt.xticks(x, [row["checkpoint"] for row in expert_rows], rotation=20, ha="right")
    plt.ylabel("Mean score diff vs heuristic_expert")
    plt.title("External anchor trajectory")
    plt.tight_layout()
    plt.savefig(paths["expert"], dpi=160)
    plt.close()

    plt.figure(figsize=(8.2, 4.4))
    protocols = sorted({row["protocol"] for row in exploiter_rows})
    for protocol in protocols:
        series = [row for row in exploiter_rows if row["protocol"] == protocol]
        plt.plot(
            [row["target"] for row in series],
            [row["win_rate"] for row in series],
            "o-",
            label=protocol,
        )
    plt.axhline(0.55, color="black", linewidth=1, linestyle="--", label="0.55 gate")
    plt.xticks(rotation=20, ha="right")
    plt.ylabel("Exploiter win rate")
    plt.title("Exploitability by protocol")
    plt.legend(fontsize=8)
    plt.tight_layout()
    plt.savefig(paths["exploiter"], dpi=160)
    plt.close()

    return {key: str(path) for key, path in paths.items()}


def build_report(
    h2h_rows: list[dict[str, Any]],
    expert_rows: list[dict[str, Any]],
    exploiter_rows: list[dict[str, Any]],
    plot_paths: dict[str, str],
) -> str:
    expert_recent = expert_rows[-1]["mean_score_diff"] - expert_rows[-2]["mean_score_diff"]
    expert_recent_ci_crosses_zero = (
        expert_rows[-1]["score_diff_ci95_low"] <= expert_rows[-2]["score_diff_ci95_high"]
        and expert_rows[-2]["score_diff_ci95_low"] <= expert_rows[-1]["score_diff_ci95_high"]
    )
    same_protocol_repairs = [
        row
        for row in exploiter_rows
        if row["phase"] == "repair_c01" and row["delta_from_previous_same_protocol"] is not None
    ]
    repair_deltas = [row["delta_from_previous_same_protocol"] for row in same_protocol_repairs]
    mean_repair_drop = -sum(repair_deltas) / len(repair_deltas)
    expected = estimate_expected_improvement(h2h_rows, expert_recent, mean_repair_drop)
    judgment = (
        f"추가 학습 1사이클(~1.5시간 GPU)의 기대 개선은 인접 H2H 최근 이득 "
        f"{h2h_rows[-1]['mean_score_diff']:+.2f}점, expert 앵커 최근 변화 "
        f"{expert_recent:+.2f}점, 동일 프로토콜 exploiter 평균 피탈률 감소 "
        f"{mean_repair_drop:.3f}에 따라 대략 {expected} 수준으로 추정되며, "
        f"수확체감 구간에 진입했다. 근거는 최근 H2H 이득이 첫 리그 전이보다 작고, "
        f"expert 앵커 성능이 v1 update 250 이후 하락/회복을 반복하며, 보수 후 "
        f"exploiter 최악값이 0.549로 게이트 바로 아래에 머문다는 세 측정의 일치다."
    )
    if not expert_recent_ci_crosses_zero:
        judgment += " 단, repair가 expert 앵커를 유의하게 회복시킨 점은 별도 긍정 신호다."

    lines = [
        f"# Diminishing Returns Diagnostic - {DATE}",
        "",
        "신규 학습 없이 기존 체크포인트와 평가 롤아웃만 사용했다. 모든 duplicate 평가는 "
        f"셔플 뱅크 seed 20260704, 2,000쌍({TOTAL_DUPLICATE_GAMES} games) 기준이다.",
        "",
        f"Raw artifacts: `{ARTIFACT_DIR}`",
        "",
        "## Checkpoints",
        "",
        "| Order | Name | Config | Checkpoint |",
        "| ---: | --- | --- | --- |",
    ]
    for spec in CHECKPOINTS:
        lines.append(
            f"| {spec.stage_index} | `{spec.name}` | `{spec.config}` | `{spec.checkpoint}` |"
        )

    lines.extend(
        [
            "",
            "## 1. Adjacent Head-to-Head",
            "",
            f"![Adjacent H2H]({Path(plot_paths['h2h']).name})",
            "",
            "| Later | Earlier | Win rate | Wilson CI | Mean diff | Score CI | Opened colors | Max-step |",
            "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |",
        ]
    )
    for row in h2h_rows:
        lines.append(match_table_row(row, row["current"], row["previous"]))

    lines.extend(
        [
            "",
            "## 2. External Anchor Trajectory",
            "",
            f"![Expert anchor]({Path(plot_paths['expert']).name})",
            "",
            "| Checkpoint | Win rate | Wilson CI | Mean diff | Score CI | Opened colors | Max-step |",
            "| --- | ---: | ---: | ---: | ---: | ---: | ---: |",
        ]
    )
    for row in expert_rows:
        lines.append(anchor_table_row(row))
    lines.extend(
        [
            "",
            f"Recent expert-anchor slope: `{expert_recent:+.4f}` points "
            f"({expert_rows[-2]['checkpoint']} -> {expert_rows[-1]['checkpoint']}).",
            "",
            "## 3. Exploitability Trajectory",
            "",
            f"![Exploitability]({Path(plot_paths['exploiter']).name})",
            "",
            "프로토콜이 다른 피탈률은 한 곡선에 섞지 않았다. `Delta`는 같은 프로토콜 안에서만 계산했다.",
            "",
            "| Phase | Target | Protocol | Budget | Warm start | Win rate | Delta same protocol | Mean diff | Source |",
            "| --- | --- | --- | --- | --- | ---: | ---: | ---: | --- |",
        ]
    )
    for row in exploiter_rows:
        delta = row["delta_from_previous_same_protocol"]
        delta_text = "n/a" if delta is None else f"{delta:+.5f}"
        source_name = Path(row["source"]).name
        lines.append(
            f"| `{row['phase']}` | `{row['target']}` | `{row['protocol']}` | "
            f"{row['budget']} | `{row['warmstart']}` | {row['win_rate']:.5f} | "
            f"{delta_text} | {row['mean_score_diff']:+.4f} | [{source_name}]({source_name}) |"
        )

    lines.extend(
        [
            "",
            "## Judgment",
            "",
            judgment,
            "",
            "## Notes",
            "",
            "- H2H와 expert 앵커 평가는 이번 작업에서 새로 실행했다.",
            "- Exploiter 표는 지금까지 생성된 리포트/JSONL의 기존 측정값만 재정리했다.",
            "- 이 작업에서는 신규 학습, 봇 수정, 체크포인트 수정이 없었다.",
        ]
    )
    return "\n".join(lines) + "\n"


def match_table_row(row: dict[str, Any], left: str, right: str) -> str:
    return (
        f"| `{left}` | `{right}` | {row['win_rate']:.4f} | "
        f"[{row['wilson_low']:.4f}, {row['wilson_high']:.4f}] | "
        f"{row['mean_score_diff']:+.4f} | "
        f"[{row['score_diff_ci95_low']:+.4f}, {row['score_diff_ci95_high']:+.4f}] | "
        f"{row['opened_colors_per_game']:.4f} | {row['max_steps_rate']:.4f} |"
    )


def anchor_table_row(row: dict[str, Any]) -> str:
    return (
        f"| `{row['checkpoint']}` | {row['win_rate']:.4f} | "
        f"[{row['wilson_low']:.4f}, {row['wilson_high']:.4f}] | "
        f"{row['mean_score_diff']:+.4f} | "
        f"[{row['score_diff_ci95_low']:+.4f}, {row['score_diff_ci95_high']:+.4f}] | "
        f"{row['opened_colors_per_game']:.4f} | {row['max_steps_rate']:.4f} |"
    )


def estimate_expected_improvement(
    h2h_rows: list[dict[str, Any]],
    expert_recent: float,
    mean_repair_drop: float,
) -> str:
    latest_h2h = h2h_rows[-1]["mean_score_diff"]
    previous_h2h = h2h_rows[-2]["mean_score_diff"]
    conservative_h2h = max(0.0, min(latest_h2h, previous_h2h))
    expert_component = max(0.0, expert_recent)
    return (
        f"+{conservative_h2h:.1f}~+{max(conservative_h2h, expert_component):.1f}점 "
        f"또는 피탈률 -{mean_repair_drop:.3f} 내외"
    )


def write_summary(
    h2h_rows: list[dict[str, Any]],
    expert_rows: list[dict[str, Any]],
    exploiter_rows: list[dict[str, Any]],
    elapsed_seconds: float,
) -> None:
    rows = []
    rows.extend(h2h_rows)
    rows.extend(expert_rows)
    rows.extend({"event": "exploiter_trajectory", **row} for row in exploiter_rows)
    rows.append(
        {
            "event": "diminishing_returns_complete",
            "elapsed_seconds": elapsed_seconds,
            "artifact_dir": str(ARTIFACT_DIR),
            "report": str(REPORT_PATH),
        }
    )
    SUMMARY_PATH.write_text(
        "\n".join(json.dumps(row, ensure_ascii=False, sort_keys=True) for row in rows) + "\n",
        encoding="utf-8",
    )


if __name__ == "__main__":
    main()
