#!/usr/bin/env python3
"""Verify pinned Hugging Face ModernBERT -> MLX backbone parity."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Sequence

import numpy as np

from deep_contract import DEEP_VARIANTS, validate_variant_config
from purpose_data import DataError
from train_deep_mlx import _configure_mlx_device, _resolve_source


def verify(variant_name: str, local_model: Path | None, device: str) -> float:
    try:
        import mlx.core as mx
        import torch
        from transformers import ModernBertForMaskedLM

        from deep_model_mlx import (
            ModernBertForPurposeClassification,
            ModernBertPurposeConfig,
            load_pretrained_weights,
        )
    except ImportError as exc:
        raise DataError(
            "deep parity requires PyTorch, Transformers, and MLX"
        ) from exc

    _configure_mlx_device(mx, device)
    variant = DEEP_VARIANTS[variant_name]
    source = _resolve_source(variant, local_model)
    try:
        config_json = json.loads(
            (source / "config.json").read_text(encoding="utf-8")
        )
    except (OSError, json.JSONDecodeError) as exc:
        raise DataError(f"cannot read ModernBERT config: {exc}") from exc
    validate_variant_config(config_json, variant)

    torch.set_num_threads(1)
    reference = ModernBertForMaskedLM.from_pretrained(
        source,
        local_files_only=True,
        attn_implementation="eager",
    )
    reference.eval()
    candidate = ModernBertForPurposeClassification(
        ModernBertPurposeConfig.from_hugging_face(
            config_json, gradient_checkpointing=False
        )
    )
    report = load_pretrained_weights(candidate, source / "model.safetensors")
    candidate.eval()

    # 96 tokens crosses the local layer's 64-token half-window, so this catches
    # both full-attention and sliding-window-mask parity without a slow 512-token
    # CPU reference pass.
    rng = np.random.default_rng(20260731)
    input_ids = rng.integers(
        3, int(config_json["vocab_size"]) - 1, size=(2, 96), dtype=np.int32
    )
    input_ids[:, 0] = int(config_json["bos_token_id"])
    input_ids[:, -1] = int(config_json["eos_token_id"])
    attention_mask = np.ones_like(input_ids, dtype=np.int32)
    attention_mask[0, -11:] = 0
    input_ids[0, -11:] = int(config_json["pad_token_id"])

    with torch.inference_mode():
        sequence = reference.model(
            input_ids=torch.from_numpy(input_ids.astype(np.int64)),
            attention_mask=torch.from_numpy(attention_mask.astype(np.int64)),
        ).last_hidden_state
        reference_pooled = reference.head(sequence[:, 0]).cpu().numpy()
    candidate_sequence = candidate.model(
        mx.array(input_ids), mx.array(attention_mask)
    )
    candidate_pooled = candidate.head(candidate_sequence[:, 0])
    mx.eval(candidate_pooled)
    error = float(
        np.max(np.abs(reference_pooled - np.asarray(candidate_pooled)))
    )
    print(
        f"{variant.model_id}@{variant.revision}: "
        f"loaded={report['loaded']} ignored={report['ignored']} "
        f"pooled_max_abs_error={error:.3g}",
        flush=True,
    )
    if not np.isfinite(error) or error > 5e-4:
        raise DataError(
            f"ModernBERT MLX parity failed: max abs error {error:.6g} > 0.0005"
        )
    return error


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--variant", choices=tuple(DEEP_VARIANTS), default="base")
    parser.add_argument("--model", type=Path)
    parser.add_argument("--device", choices=("metal", "cpu"), default="metal")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        verify(args.variant, args.model, args.device)
    except (DataError, OSError, RuntimeError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


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