#!/usr/bin/env python3
"""Compare legacy and safetensors model predictions on one prepared sample."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
from pathlib import Path

import torch
import yaml


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while block := stream.read(16 * 1024 * 1024):
            digest.update(block)
    return digest.hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-root", required=True, type=Path)
    parser.add_argument("--legacy-checkpoint", required=True, type=Path)
    parser.add_argument("--release-checkpoint", required=True, type=Path)
    parser.add_argument("--model-config", required=True, type=Path)
    parser.add_argument("--timestamp", required=True)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--output", type=Path, default=None)
    args = parser.parse_args()

    modeling_root = Path(__file__).resolve().parents[2] / "modeling"
    if str(modeling_root) not in sys.path:
        sys.path.insert(0, str(modeling_root))
    from src.data_pipeline.fast_dataset import FastDataset
    from src.training_validation.common import build_model, load_model_checkpoint

    requested = args.device
    if requested == "auto":
        requested = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(requested)
    with args.model_config.open("r", encoding="utf-8") as stream:
        model_config = yaml.safe_load(stream)

    dataset_config = {
        "output_root": str(args.dataset_root.resolve()),
        "catalog_path": str((args.dataset_root / "catalog.csv").resolve()),
        "input_window": {"past_minutes": 50, "interval_minutes": 10},
        "splits": {"check": [{"start": args.timestamp, "end": args.timestamp}]},
    }
    dataset = FastDataset(
        dataset_config,
        split="check",
        required_inputs=["concat", "hsr"],
        required_labels=[],
    )
    if dataset.samples != [args.timestamp]:
        raise ValueError(f"timestamp is not a complete inference sample: {args.timestamp}")
    item = dataset[0]
    x = torch.cat([item["inputs"][name] for name in ("concat", "hsr")], dim=1).unsqueeze(0).to(device)

    model = build_model(model_config).to(device).eval()
    with torch.inference_mode():
        load_model_checkpoint(model, args.legacy_checkpoint, device)
        legacy = {name: value.detach().cpu() for name, value in model(x).items()}
        load_model_checkpoint(model, args.release_checkpoint, device)
        released = {name: value.detach().cpu() for name, value in model(x).items()}

    errors = {
        name: float(torch.max(torch.abs(legacy[name] - released[name])).item())
        for name in legacy
    }
    result = {
        "timestamp": args.timestamp,
        "device": str(device),
        "legacy_checkpoint": {"filename": args.legacy_checkpoint.name, "sha256": sha256(args.legacy_checkpoint)},
        "release_checkpoint": {"filename": args.release_checkpoint.name, "sha256": sha256(args.release_checkpoint)},
        "max_absolute_error": errors,
        "tolerance": 1e-5,
        "passed": all(value <= 1e-5 for value in errors.values()),
    }
    if not result["passed"]:
        raise AssertionError(f"prediction mismatch: {errors}")
    output = args.output or args.release_checkpoint.parent / "prediction_equivalence.json"
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_name(f".{output.name}.partial")
    with temporary.open("w", encoding="utf-8") as stream:
        json.dump(result, stream, indent=2)
        stream.write("\n")
    os.replace(temporary, output)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()