#!/usr/bin/env python3
"""Verify MLX/PyTorch parity and checkpoint round-tripping before MLX training."""

from __future__ import annotations

import argparse
import tempfile
from pathlib import Path
from typing import Sequence

import numpy as np

from purpose_data import LABELS, DataError, load_jsonl
from train import enable_quantization_aware_training, encode_fixed_shape
from train_mlx import (
    _checkpoint_config,
    _configure_mlx_device,
    encode_fixed_shape_numpy,
)


SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_DATASET = SCRIPT_DIR / ".artifacts" / "dataset-v1" / "validation.jsonl"


def verify(model_dir: Path, dataset: Path, records: int, device: str) -> None:
    try:
        import mlx.core as mx
        import mlx.nn as nn
        import torch
        from mlx.utils import tree_flatten
        from safetensors import safe_open
        from transformers import AutoModelForSequenceClassification, AutoTokenizer

        from mlx_model import (
            BertClassifierConfig,
            BertForSequenceClassification,
            fake_quantize_activation,
            fake_quantize_linear_weight,
            load_hugging_face_weights,
            save_hugging_face_weights,
        )
    except ImportError as exc:
        raise DataError(
            "verification requires requirements-mlx.txt"
        ) from exc
    _configure_mlx_device(mx, device)

    # Check the fake-quantization contract independently of the full model. Tiny
    # backend-specific floating-point differences can cross later quantization
    # thresholds, especially on padded tokens, so full QAT logits are expected to
    # have more drift than the ordinary float graph.
    activation_values = np.asarray(
        [[-2.75, -0.125, 0.0, 0.625], [1.25, 3.5, -1.0, 0.25]],
        dtype=np.float32,
    )
    torch_activation = torch.from_numpy(activation_values)
    activation_minimum = min(0.0, float(torch_activation.amin().item()))
    activation_maximum = max(0.0, float(torch_activation.amax().item()))
    activation_scale = max(
        (activation_maximum - activation_minimum) / 255.0,
        torch.finfo(torch.float32).eps,
    )
    activation_zero_point = max(
        0,
        min(255, round(-activation_minimum / activation_scale)),
    )
    torch_quantized_activation = torch.fake_quantize_per_tensor_affine(
        torch_activation,
        activation_scale,
        activation_zero_point,
        0,
        255,
    ).numpy()
    mlx_quantized_activation = np.asarray(
        fake_quantize_activation(mx.array(activation_values))
    )
    np.testing.assert_array_equal(
        mlx_quantized_activation,
        torch_quantized_activation,
    )

    weight_values = np.asarray(
        [[-1.5, -0.25, 0.75, 1.25], [0.125, -0.875, 2.0, -1.25]],
        dtype=np.float32,
    )
    torch_weight = torch.from_numpy(weight_values)
    weight_scales = torch_weight.abs().amax(dim=1).div(127.0).clamp_min(
        torch.finfo(torch.float32).eps
    )
    torch_quantized_weight = torch.fake_quantize_per_channel_affine(
        torch_weight,
        weight_scales,
        torch.zeros_like(weight_scales, dtype=torch.int32),
        0,
        -127,
        127,
    ).numpy()
    mlx_quantized_weight = np.asarray(
        fake_quantize_linear_weight(mx.array(weight_values))
    )
    np.testing.assert_array_equal(mlx_quantized_weight, torch_quantized_weight)
    print("fake-quant primitives: exact")

    checkpoint = model_dir / "model.safetensors"
    if not checkpoint.is_file():
        raise DataError(f"{checkpoint}: checkpoint is missing")
    validation = load_jsonl(dataset)[:records]
    if not validation:
        raise DataError(f"{dataset}: no validation records")
    texts = [record["prompt"] for record in validation]
    labels = np.asarray(
        [LABELS.index(record["purpose"]) for record in validation],
        dtype=np.int32,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
    numpy_tokens = encode_fixed_shape_numpy(tokenizer, texts)
    torch_tokens = encode_fixed_shape(tokenizer, texts, torch)

    config_json = _checkpoint_config(model_dir)
    config = BertClassifierConfig.from_hugging_face(config_json)
    mlx_float = BertForSequenceClassification(
        config,
        quantization_aware=False,
    )
    load_hugging_face_weights(mlx_float, checkpoint)
    mlx_float.eval()
    mlx_float_logits = mlx_float(
        **{key: mx.array(value) for key, value in numpy_tokens.items()}
    )
    mx.eval(mlx_float_logits)
    mlx_float_logits = np.asarray(mlx_float_logits)

    torch_model = AutoModelForSequenceClassification.from_pretrained(
        model_dir,
        local_files_only=True,
    )
    torch_model.eval()
    with torch.inference_mode():
        torch_float_logits = torch_model(**torch_tokens).logits.numpy()
    np.testing.assert_allclose(
        mlx_float_logits,
        torch_float_logits,
        rtol=1e-4,
        atol=2e-5,
    )
    if not np.array_equal(
        mlx_float_logits.argmax(axis=-1),
        torch_float_logits.argmax(axis=-1),
    ):
        raise DataError("MLX and PyTorch float predictions differ")
    print(
        "float parity: "
        f"max_abs_error={np.max(np.abs(mlx_float_logits - torch_float_logits)):.3g}"
    )

    with tempfile.TemporaryDirectory(prefix="purpose-mlx-roundtrip-") as temp:
        roundtrip = Path(temp) / "model.safetensors"
        save_hugging_face_weights(mlx_float, roundtrip)
        with safe_open(checkpoint, framework="np") as original, safe_open(
            roundtrip,
            framework="np",
        ) as converted:
            if set(original.keys()) != set(converted.keys()):
                raise DataError("MLX checkpoint round-trip changed parameter keys")
            for key in original.keys():
                np.testing.assert_array_equal(
                    original.get_tensor(key),
                    converted.get_tensor(key),
                )
    print("checkpoint round-trip: exact")

    del mlx_float
    mx.clear_cache()
    mlx_qat = BertForSequenceClassification(
        config,
        quantization_aware=True,
    )
    load_hugging_face_weights(mlx_qat, checkpoint)
    mlx_qat.eval()
    mlx_qat_logits = mlx_qat(
        **{key: mx.array(value) for key, value in numpy_tokens.items()}
    )
    mx.eval(mlx_qat_logits)
    mlx_qat_logits = np.asarray(mlx_qat_logits)

    enable_quantization_aware_training(torch, torch_model)
    torch_model.eval()
    with torch.inference_mode():
        torch_qat_logits = torch_model(**torch_tokens).logits.numpy()
    qat_max_abs_error = float(
        np.max(np.abs(mlx_qat_logits - torch_qat_logits))
    )
    if not np.isfinite(qat_max_abs_error) or qat_max_abs_error > 1.0:
        raise DataError(
            "MLX and PyTorch QAT logits have excessive backend drift: "
            f"{qat_max_abs_error:.3g}"
        )
    if not np.array_equal(
        mlx_qat_logits.argmax(axis=-1),
        torch_qat_logits.argmax(axis=-1),
    ):
        raise DataError("MLX and PyTorch QAT predictions differ")
    print(
        "QAT parity: "
        f"predictions=exact max_abs_error={qat_max_abs_error:.3g}"
    )

    mlx_qat.train()
    mlx_labels = mx.array(labels)

    def loss_function() -> object:
        logits = mlx_qat(
            **{key: mx.array(value) for key, value in numpy_tokens.items()}
        )
        return nn.losses.cross_entropy(logits, mlx_labels, reduction="mean")

    loss, gradients = nn.value_and_grad(mlx_qat, loss_function)()
    flat_gradients = [gradient for _, gradient in tree_flatten(gradients)]
    mx.eval(loss, gradients)
    if not all(
        bool(mx.all(mx.isfinite(gradient)).item())
        for gradient in flat_gradients
    ):
        raise DataError("MLX QAT produced non-finite gradients")
    if not any(
        float(mx.max(mx.abs(gradient)).item()) > 0
        for gradient in flat_gradients
    ):
        raise DataError("MLX QAT produced only zero gradients")
    print(f"QAT gradient smoke: loss={float(loss.item()):.6f}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--dataset", type=Path, default=DEFAULT_DATASET)
    parser.add_argument("--records", type=int, default=8)
    parser.add_argument(
        "--device",
        choices=("metal", "cpu"),
        default="metal",
        help="MLX execution device (default: metal; cpu is a diagnostic fallback)",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if args.records <= 0:
        raise SystemExit("--records must be positive")
    try:
        verify(
            args.model.expanduser(),
            args.dataset.expanduser(),
            args.records,
            args.device,
        )
    except (AssertionError, DataError) as exc:
        print(f"error: {exc}")
        return 2
    return 0


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