#!/usr/bin/env python3
"""Compare Tinker checkpoints on the fixed hidden-variant tool environment."""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import logging
import math
import re
import sys
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import tinker

EXAMPLE_DIR = Path(__file__).resolve().parent
SRC_DIR = EXAMPLE_DIR / "src"
sys.path.insert(0, str(SRC_DIR))

from euler_tool_rl import (  # noqa: E402
    build_variant_env,
    close_sandbox_clients,
    verify_sandbox_languages,
)
from euler_variants import (  # noqa: E402
    FAMILIES,
    LANGUAGES,
    VariantTask,
    assert_split_isolation,
    build_variant_tasks,
    ensure_loopback_sandbox_url,
    task_set_digest,
)
from tinker_cookbook import renderers  # noqa: E402
from tinker_cookbook.eval.benchmarks import (  # noqa: E402
    BenchmarkBuilder,
    BenchmarkConfig,
    BenchmarkResult,
    run_benchmark,
)
from tinker_cookbook.rl.types import Env  # noqa: E402
from tinker_cookbook.tokenizer_utils import get_tokenizer  # noqa: E402

BENCHMARK_NAME = "project_euler_hidden_variants"
CHECKPOINT_LABEL = re.compile(r"^[a-z][a-z0-9_-]{0,31}$")
DEFAULT_EVAL_DIGEST = "31277b8e51831ce17722e5e3bb40898fb3f3b4f32549fa353c1be5361569a30d"


def comma_values(spec: str) -> list[str]:
    values = [value.strip() for value in spec.split(",") if value.strip()]
    if not values:
        raise ValueError("at least one value is required")
    return values


def checkpoint_spec(spec: str) -> tuple[str, str]:
    label, separator, path = spec.partition("=")
    if not separator or not CHECKPOINT_LABEL.fullmatch(label):
        raise ValueError(
            "checkpoint must use label=tinker://path with a simple lowercase label"
        )
    if not path.startswith("tinker://"):
        raise ValueError(f"checkpoint {label} must use a tinker:// path")
    return label, path


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def wilson_interval(
    successes: int, total: int, z: float = 1.959963984540054
) -> list[float]:
    if total == 0:
        return [0.0, 0.0]
    rate = successes / total
    denominator = 1.0 + z * z / total
    center = (rate + z * z / (2 * total)) / denominator
    radius = (
        z
        * math.sqrt(rate * (1 - rate) / total + z * z / (4 * total * total))
        / denominator
    )
    return [max(0.0, center - radius), min(1.0, center + radius)]


def fisher_exact_two_sided(
    success_a: int, total_a: int, success_b: int, total_b: int
) -> float:
    row_a = total_a
    row_b = total_b
    column_success = success_a + success_b
    total = row_a + row_b
    minimum = max(0, column_success - row_b)
    maximum = min(row_a, column_success)

    def probability(cell_a: int) -> float:
        return (
            math.comb(column_success, cell_a)
            * math.comb(total - column_success, row_a - cell_a)
            / math.comb(total, row_a)
        )

    observed = probability(success_a)
    return min(
        1.0,
        sum(
            probability(cell_a)
            for cell_a in range(minimum, maximum + 1)
            if probability(cell_a) <= observed + 1e-15
        ),
    )


def rate_difference(
    source: dict[str, Any], candidate: dict[str, Any]
) -> dict[str, Any]:
    source_successes = int(source["full_passes"])
    source_total = int(source["samples"])
    candidate_successes = int(candidate["full_passes"])
    candidate_total = int(candidate["samples"])
    source_rate = source_successes / source_total if source_total else 0.0
    candidate_rate = candidate_successes / candidate_total if candidate_total else 0.0
    difference = candidate_rate - source_rate
    standard_error = math.sqrt(
        source_rate * (1.0 - source_rate) / source_total
        + candidate_rate * (1.0 - candidate_rate) / candidate_total
    )
    z = 1.959963984540054
    return {
        "method_note": (
            "Descriptive unpaired trajectory-level approximation. Conditions are fixed "
            "and repeated; this is not a cluster-robust or paired causal estimate."
        ),
        "source_full_passes": source_successes,
        "source_samples": source_total,
        "candidate_full_passes": candidate_successes,
        "candidate_samples": candidate_total,
        "full_pass_rate_delta": difference,
        "unpooled_normal_difference_95": [
            difference - z * standard_error,
            difference + z * standard_error,
        ],
        "fisher_exact_two_sided_p": fisher_exact_two_sided(
            candidate_successes,
            candidate_total,
            source_successes,
            source_total,
        ),
    }


@dataclass
class HiddenVariantBenchmark(BenchmarkBuilder):
    tasks: list[VariantTask]
    languages: list[str]
    model_name: str
    sandbox_url: str
    sandbox_timeout: int
    max_turns: int
    max_tool_calls: int
    max_trajectory_tokens: int
    max_generation_tokens: int
    name: str = BENCHMARK_NAME
    multi_turn: bool = True
    recommended_timeout: float = 180.0

    def make_envs(
        self,
        renderer: renderers.Renderer,
        _config: BenchmarkConfig,
    ) -> Sequence[Env]:
        return [
            build_variant_env(
                task=task,
                language=language,
                model_name=self.model_name,
                renderer=renderer,
                sandbox_url=self.sandbox_url,
                sandbox_timeout=self.sandbox_timeout,
                max_turns=self.max_turns,
                max_tool_calls=self.max_tool_calls,
                max_trajectory_tokens=self.max_trajectory_tokens,
                max_generation_tokens=self.max_generation_tokens,
                evaluation_reward=True,
                example_id=f"{task.task_id}:{language}",
            )
            for task in self.tasks
            for language in self.languages
        ]

    def aggregate(
        self,
        rewards: list[float],
        metrics_list: list[dict[str, float | int]],
    ) -> BenchmarkResult:
        full_passes = sum(
            int(metrics.get("hidden_all_passed", 0)) for metrics in metrics_list
        )
        count = len(rewards)

        def mean_metric(name: str) -> float:
            return (
                sum(float(metrics.get(name, 0.0)) for metrics in metrics_list) / count
                if count
                else 0.0
            )

        return BenchmarkResult(
            name=self.name,
            score=full_passes / count if count else 0.0,
            num_examples=count,
            num_correct=full_passes,
            metrics={
                "hidden_case_pass_rate": mean_metric("hidden_pass_rate"),
                "tool_call_rate": mean_metric("tool_called"),
                "submission_error_rate": mean_metric("submission_error"),
                "infrastructure_error_rate": mean_metric("infrastructure_error"),
            },
        )


def read_trajectories(path: Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                records.append(json.loads(line))
    return records


def summarize_records(
    records: list[dict[str, Any]],
    tasks: list[VariantTask],
    families: list[str],
    languages: list[str],
    result: BenchmarkResult,
) -> dict[str, Any]:
    task_family = {task.task_id: task.family for task in tasks}
    full_passes_by_example: dict[str, list[int]] = defaultdict(list)
    buckets: dict[str, dict[str, float]] = defaultdict(
        lambda: {
            "samples": 0.0,
            "full_passes": 0.0,
            "hidden_case_pass_rate_sum": 0.0,
            "tool_calls": 0.0,
            "submission_errors": 0.0,
            "infrastructure_errors": 0.0,
            "action_tokens": 0.0,
            "observation_tokens": 0.0,
        }
    )

    for record in records:
        example_id = str(record["example_id"])
        task_id, language = example_id.rsplit(":", 1)
        family = task_family[task_id]
        metrics = record.get("metrics", {})
        full_passes_by_example[example_id].append(
            int(metrics.get("hidden_all_passed", 0.0))
        )
        action_tokens = sum(
            int(turn.get("token_count", 0))
            for turn in record.get("turns", [])
            if turn.get("role") == "assistant"
        )
        observation_tokens = sum(
            int(turn.get("token_count", 0))
            for turn in record.get("turns", [])
            if turn.get("role") != "assistant"
        )
        for key in (
            "all",
            f"language:{language}",
            f"family:{family}",
            f"condition:{family}:{language}",
        ):
            bucket = buckets[key]
            bucket["samples"] += 1
            bucket["full_passes"] += float(metrics.get("hidden_all_passed", 0.0))
            bucket["hidden_case_pass_rate_sum"] += float(
                metrics.get("hidden_pass_rate", 0.0)
            )
            bucket["tool_calls"] += float(metrics.get("tool_called", 0.0))
            bucket["submission_errors"] += float(metrics.get("submission_error", 0.0))
            bucket["infrastructure_errors"] += float(
                metrics.get("infrastructure_error", 0.0)
            )
            bucket["action_tokens"] += action_tokens
            bucket["observation_tokens"] += observation_tokens

    def finish(bucket: dict[str, float]) -> dict[str, Any]:
        samples = int(bucket["samples"])
        full_passes = int(bucket["full_passes"])
        return {
            "samples": samples,
            "full_passes": full_passes,
            "full_pass_rate": full_passes / samples if samples else 0.0,
            "full_pass_wilson_95": wilson_interval(full_passes, samples),
            "hidden_case_pass_rate": (
                bucket["hidden_case_pass_rate_sum"] / samples if samples else 0.0
            ),
            "tool_call_rate": bucket["tool_calls"] / samples if samples else 0.0,
            "submission_error_rate": (
                bucket["submission_errors"] / samples if samples else 0.0
            ),
            "infrastructure_error_rate": (
                bucket["infrastructure_errors"] / samples if samples else 0.0
            ),
            "action_tokens": int(bucket["action_tokens"]),
            "observation_tokens": int(bucket["observation_tokens"]),
        }

    binary_pass_at_k: dict[str, float] = {}
    for key, runner_value in result.pass_at_k.items():
        k = int(key)
        estimates: list[float] = []
        for outcomes in full_passes_by_example.values():
            samples = len(outcomes)
            successes = sum(outcomes)
            if k > samples:
                continue
            failure_combinations = (
                math.comb(samples - successes, k) if samples - successes >= k else 0
            )
            estimates.append(1.0 - failure_combinations / math.comb(samples, k))
        value = sum(estimates) / len(estimates) if estimates else 0.0
        if not math.isclose(value, float(runner_value), rel_tol=0.0, abs_tol=1e-12):
            raise RuntimeError(
                f"binary pass@{k} mismatch: recomputed {value}, runner {runner_value}"
            )
        binary_pass_at_k[str(k)] = value

    return {
        "overall": finish(buckets["all"]),
        "by_language": {
            language: finish(buckets[f"language:{language}"]) for language in languages
        },
        "by_family": {
            family: finish(buckets[f"family:{family}"]) for family in families
        },
        "by_condition": {
            f"{family}:{language}": finish(buckets[f"condition:{family}:{language}"])
            for family in families
            for language in languages
        },
        "pass_at_k": binary_pass_at_k,
        "runner_errors": result.num_errors,
        "runner_truncations": result.num_truncated,
        "elapsed_seconds": result.time_seconds,
    }


def comparison_summary(
    summaries: dict[str, dict[str, Any]], labels: list[str] | None = None
) -> dict[str, Any]:
    labels = labels or list(summaries)
    if len(labels) != 2:
        return {"labels": labels, "pairwise": None}
    if set(labels) == {"source", "candidate"}:
        source_label, candidate_label = "source", "candidate"
    else:
        source_label, candidate_label = labels
    source = summaries[source_label]
    candidate = summaries[candidate_label]
    languages = list(source["by_language"])
    families = list(source["by_family"])

    return {
        "labels": labels,
        "pairwise": {
            "source": source_label,
            "candidate": candidate_label,
            "overall_inference": rate_difference(
                source["overall"], candidate["overall"]
            ),
            "hidden_case_pass_rate_delta": (
                candidate["overall"]["hidden_case_pass_rate"]
                - source["overall"]["hidden_case_pass_rate"]
            ),
            "by_language_inference": {
                language: rate_difference(
                    source["by_language"][language],
                    candidate["by_language"][language],
                )
                for language in languages
            },
            "by_family_inference": {
                family: rate_difference(
                    source["by_family"][family],
                    candidate["by_family"][family],
                )
                for family in families
            },
        },
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--checkpoint",
        action="append",
        required=True,
        help="Repeat label=tinker://sampler_path for each checkpoint",
    )
    parser.add_argument("--model", default="Qwen/Qwen3.5-4B")
    parser.add_argument("--renderer", default="qwen3_5_disable_thinking")
    parser.add_argument("--families", default=",".join(FAMILIES))
    parser.add_argument("--languages", default=",".join(LANGUAGES))
    parser.add_argument("--tasks-per-family", type=int, default=1)
    parser.add_argument("--cases-per-task", type=int, default=4)
    parser.add_argument("--seed", type=int, default=1729)
    parser.add_argument(
        "--expected-eval-digest",
        default=DEFAULT_EVAL_DIGEST,
        help="Required SHA-256 digest of the frozen evaluation task set",
    )
    parser.add_argument("--fanout", type=int, default=8)
    parser.add_argument("--max-tokens", type=int, default=2048)
    parser.add_argument("--context-window", type=int, default=32768)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--concurrency", type=int, default=6)
    parser.add_argument("--timeout", type=float, default=180.0)
    parser.add_argument("--max-turns", type=int, default=3)
    parser.add_argument("--max-tool-calls", type=int, default=2)
    parser.add_argument("--sandbox-url", default="http://127.0.0.1:8080/run_code")
    parser.add_argument("--sandbox-timeout", type=int, default=10)
    parser.add_argument("--output-dir", type=Path)
    return parser


async def run(args: argparse.Namespace) -> dict[str, Any]:
    checkpoints = [checkpoint_spec(spec) for spec in args.checkpoint]
    if len({label for label, _ in checkpoints}) != len(checkpoints):
        raise ValueError("checkpoint labels must be unique")
    if args.fanout <= 0:
        raise ValueError("--fanout must be positive")
    families = comma_values(args.families)
    languages = comma_values(args.languages)
    unknown_families = sorted(set(families) - set(FAMILIES))
    unknown_languages = sorted(set(languages) - set(LANGUAGES))
    if unknown_families:
        raise ValueError(f"unknown families: {', '.join(unknown_families)}")
    if unknown_languages:
        raise ValueError(f"unknown languages: {', '.join(unknown_languages)}")

    train_tasks = build_variant_tasks(
        families,
        split="train",
        tasks_per_family=args.tasks_per_family,
        cases_per_task=args.cases_per_task,
        seed=args.seed,
    )
    eval_tasks = build_variant_tasks(
        families,
        split="eval",
        tasks_per_family=args.tasks_per_family,
        cases_per_task=args.cases_per_task,
        seed=args.seed,
    )
    assert_split_isolation(train_tasks, eval_tasks)
    train_digest = task_set_digest(train_tasks)
    eval_digest = task_set_digest(eval_tasks)
    if eval_digest != args.expected_eval_digest:
        raise RuntimeError(
            "evaluation task digest mismatch: "
            f"expected {args.expected_eval_digest}, generated {eval_digest}"
        )

    ensure_loopback_sandbox_url(args.sandbox_url)
    await verify_sandbox_languages(args.sandbox_url, args.sandbox_timeout)

    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    output_dir = (
        args.output_dir
        or EXAMPLE_DIR / "runs" / f"tool-rl-checkpoint-comparison-{timestamp}"
    ).resolve()
    if output_dir.exists():
        raise FileExistsError(f"output directory already exists: {output_dir}")
    output_dir.mkdir(parents=True)

    manifest_path = output_dir / "manifest.json"
    manifest: dict[str, Any] = {
        "schema": "machine.project-euler.tool-rl-comparison-manifest.v1",
        "status": "prepared",
        "read_only_sampling": True,
        "optimizer_created": False,
        "checkpoints": dict(checkpoints),
        "model": args.model,
        "renderer": args.renderer,
        "families": families,
        "languages": languages,
        "tasks_per_family": args.tasks_per_family,
        "cases_per_task": args.cases_per_task,
        "seed": args.seed,
        "fanout": args.fanout,
        "max_tokens": args.max_tokens,
        "temperature": args.temperature,
        "train_digest": train_digest,
        "eval_digest": eval_digest,
        "expected_eval_digest": args.expected_eval_digest,
        "source_sha256": {
            "evaluate_tool_rl.py": sha256_file(Path(__file__)),
            "src/euler_tool_rl.py": sha256_file(SRC_DIR / "euler_tool_rl.py"),
            "src/euler_variants.py": sha256_file(SRC_DIR / "euler_variants.py"),
        },
    }
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")

    renderer = renderers.get_renderer(args.renderer, get_tokenizer(args.model))
    benchmark = HiddenVariantBenchmark(
        tasks=eval_tasks,
        languages=languages,
        model_name=args.model,
        sandbox_url=args.sandbox_url,
        sandbox_timeout=args.sandbox_timeout,
        max_turns=args.max_turns,
        max_tool_calls=args.max_tool_calls,
        max_trajectory_tokens=args.context_window,
        max_generation_tokens=args.max_tokens,
    )
    service_client = tinker.ServiceClient()
    summaries: dict[str, dict[str, Any]] = {}
    manifest["status"] = "sampling"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")

    try:
        for label, checkpoint_path in checkpoints:
            checkpoint_dir = output_dir / label
            sampling_client = service_client.create_sampling_client(
                model_path=checkpoint_path
            )
            result = await run_benchmark(
                benchmark,
                sampling_client,
                renderer,
                BenchmarkConfig(
                    num_samples=args.fanout,
                    save_dir=str(checkpoint_dir),
                    agent_concurrency=args.concurrency,
                    timeout_seconds=args.timeout,
                    max_tokens=args.max_tokens,
                    max_generation_tokens=args.max_tokens,
                    max_trajectory_tokens=args.context_window,
                    context_window=args.context_window,
                    temperature=args.temperature,
                ),
            )
            trajectories_path = checkpoint_dir / BENCHMARK_NAME / "trajectories.jsonl"
            summary = summarize_records(
                read_trajectories(trajectories_path),
                eval_tasks,
                families,
                languages,
                result,
            )
            summary.update(
                {
                    "schema": "machine.project-euler.tool-rl-checkpoint-summary.v1",
                    "label": label,
                    "checkpoint": checkpoint_path,
                    "eval_digest": eval_digest,
                    "fanout": args.fanout,
                }
            )
            summaries[label] = summary
            (checkpoint_dir / "summary.json").write_text(
                json.dumps(summary, indent=2, sort_keys=True) + "\n"
            )

        comparison = {
            "schema": "machine.project-euler.tool-rl-checkpoint-comparison.v1",
            "status": "completed",
            "interpretation": (
                "Read-only stochastic checkpoint comparison on a fixed internal split. "
                "This is not an external generalization or promotion result."
            ),
            "eval_digest": eval_digest,
            "fanout": args.fanout,
            "summaries": summaries,
            **comparison_summary(summaries),
        }
        (output_dir / "comparison.json").write_text(
            json.dumps(comparison, indent=2, sort_keys=True) + "\n"
        )
        manifest["status"] = "completed"
        manifest["comparison_path"] = str(output_dir / "comparison.json")
        manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
        print(json.dumps(comparison, indent=2, sort_keys=True))
        return comparison
    except BaseException as error:
        manifest["status"] = "failed"
        manifest["error_type"] = type(error).__name__
        manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
        raise
    finally:
        await close_sandbox_clients()


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    asyncio.run(run(build_parser().parse_args()))


if __name__ == "__main__":
    main()