#!/usr/bin/env python3
"""Comparer 10 étapes continues à 5 + 5 dans des processus CPU séparés."""

import argparse
from datetime import datetime, timezone
import json
import math
from pathlib import Path
import subprocess
import sys

sys.dont_write_bytecode = True

import torch

from train import ROOT, VERSION, environment, load_checkpoint, sha256, source_identity

ATOL = 1e-12
RTOL = 0.0


def execute(output, steps, resume=None, negative=False):
    command = [sys.executable, "-B", str(ROOT / "train.py"), "--steps", str(steps), "--output", str(output)]
    if resume is not None:
        command.extend(["--resume", str(resume)])
    if negative:
        command.append("--omit-rng-restore")
    # Chaque invocation démarre un nouvel interpréteur, y compris la reprise.
    completed = subprocess.run(command, capture_output=True, text=True, encoding="utf-8", timeout=120)
    if completed.returncode != 0:
        raise ValueError("Une exécution enfant a échoué : aucune équivalence de reprise n’est validée.")
    summary = json.loads((output / "summary.json").read_text(encoding="utf-8"))
    return summary, load_checkpoint(output / "checkpoint.pt")


def difference(left, right):
    """Écart maximum récursif ; structure et entiers doivent aussi correspondre."""
    if isinstance(left, torch.Tensor) and isinstance(right, torch.Tensor):
        if left.shape != right.shape or left.dtype != right.dtype:
            return math.inf
        return float((left - right).abs().max()) if left.numel() else 0.0
    if isinstance(left, dict) and isinstance(right, dict):
        if left.keys() != right.keys():
            return math.inf
        return max((difference(left[key], right[key]) for key in left), default=0.0)
    if isinstance(left, (list, tuple)) and isinstance(right, (list, tuple)):
        if len(left) != len(right):
            return math.inf
        return max((difference(a, b) for a, b in zip(left, right)), default=0.0)
    if isinstance(left, float) and isinstance(right, (int, float)):
        return abs(left - right) if math.isfinite(left) and math.isfinite(right) else math.inf
    return 0.0 if left == right else math.inf


def compare(reference, candidate):
    numeric = {"weights_max_abs": difference(reference["model"], candidate["model"]),
               "optimizer_max_abs": difference(reference["optimizer"], candidate["optimizer"]),
               "eval_mse_abs": abs(reference["eval_mse"] - candidate["eval_mse"]),
               "loss_history_max_abs": difference([row["loss"] for row in reference["history"]], [row["loss"] for row in candidate["history"]]),
               "learning_rate_history_max_abs": difference([row["lr_used"] for row in reference["history"]], [row["lr_used"] for row in candidate["history"]]),
               "next_learning_rate_abs": abs(reference["next_lr"] - candidate["next_lr"])}
    exact = {"scheduler_equal": reference["scheduler"] == candidate["scheduler"],
             "data_progress_equal": reference["progress"] == candidate["progress"],
             "batch_order_equal": [row["batch_rows"] for row in reference["history"]] == [row["batch_rows"] for row in candidate["history"]],
             "next_rng_draws_equal": reference["next_rng_draws"] == candidate["next_rng_draws"]}
    return {**numeric, **exact, "passed": all(value <= ATOL for value in numeric.values()) and all(exact.values())}


def verify(output):
    if output.exists():
        raise ValueError("Le répertoire de preuve existe déjà. Choisissez un nouveau nom.")
    output.mkdir(parents=True, exist_ok=False)
    continuous_summary, continuous = execute(output / "continuous", 10)
    first_summary, first = execute(output / "first-half", 5)
    resumed_summary, resumed = execute(output / "resumed", 5, output / "first-half" / "checkpoint.pt")
    negative_summary, negative = execute(output / "negative-no-rng", 5, output / "first-half" / "checkpoint.pt", negative=True)
    positive = compare(continuous, resumed)
    negative_comparison = compare(continuous, negative)
    boundary = {"interruption_step": first["progress"]["step"], "interruption_epoch": first["progress"]["epoch"],
                "interruption_cursor": first["progress"]["cursor"], "dataset_rows": 24,
                "samples_seen_at_interruption": first["progress"]["samples_seen"],
                "next_rng_draws_restored": first["next_rng_draws"] == resumed_summary["resumed_next_rng_draws"],
                "next_learning_rate_restored": abs(first["next_lr"] - resumed_summary["resumed_next_lr"]) <= ATOL,
                "next_learning_rate_used": abs(first["next_lr"] - resumed["history"][5]["lr_used"]) <= ATOL,
                "next_rng_draws_negative_differ": first["next_rng_draws"] != negative_summary["resumed_next_rng_draws"]}
    passed = (positive["passed"] and not negative_comparison["passed"]
              and boundary["next_rng_draws_restored"] and boundary["next_learning_rate_restored"]
              and boundary["next_learning_rate_used"] and boundary["next_rng_draws_negative_differ"]
              and continuous["progress"]["step"] == resumed["progress"]["step"] == 10
              and continuous["progress"]["samples_seen"] == resumed["progress"]["samples_seen"] == 40)
    result = {"project": "kernova-reprise", "project_version": VERSION,
              "verified_at_utc": datetime.now(timezone.utc).isoformat(), "environment": environment(),
              "source": {**source_identity(), "verifier_sha256": sha256(ROOT / "verify_resume.py")},
              "protocol": {"device": "cpu", "continuous_steps": 10, "split_steps": [5, 5],
                           "separate_child_processes": 4, "new_process_for_resume": True,
                           "batch_size": 4, "samples_per_complete_run": 40,
                           "atol": ATOL, "rtol": RTOL, "gpu_comparison": "not_run"},
              "resume_boundary": boundary, "positive": positive,
              "negative_without_rng_restore": {**negative_comparison, "divergence_detected": not negative_comparison["passed"]},
              "observed": {"continuous_eval_mse": continuous["eval_mse"], "resumed_eval_mse": resumed["eval_mse"],
                           "negative_eval_mse": negative["eval_mse"], "next_learning_rate": resumed["next_lr"]},
              "all_checks_passed": passed}
    (output / "verification.json").write_text(json.dumps(result, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    print(json.dumps({"device": "cpu", "all_checks_passed": passed, "positive": positive["passed"],
                      "negative_divergence_detected": not negative_comparison["passed"],
                      "report": "verification.json"}))
    return 0 if passed else 1


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", required=True, type=Path, help="Nouveau dossier accueillant les quatre essais et la preuve.")
    args = parser.parse_args()
    try:
        return verify(args.output)
    except (ValueError, OSError, RuntimeError, subprocess.TimeoutExpired) as error:
        message = str(error) if isinstance(error, ValueError) else "Vérification interrompue : aucune preuve positive produite."
        print(json.dumps({"all_checks_passed": False, "error_type": type(error).__name__, "message": message}, ensure_ascii=False))
        return 2


if __name__ == "__main__":
    sys.exit(main())