#!/usr/bin/env python3
"""Bounded numerical parity check for PyTorch and native transformer inference."""

from __future__ import annotations

import argparse
import hashlib
import math
import struct
import subprocess
import sys
import tempfile
import zlib
from pathlib import Path
from typing import Sequence

import numpy as np

import train_transformer as training


PROBE_INPUT_MAGIC = b"SZTPRB01"
PROBE_OUTPUT_MAGIC = b"SZTPOU01"
PROBE_VERSION = 1
PROBE_INPUT_HEADER = struct.Struct("<8sIIIII")
PROBE_OUTPUT_HEADER = struct.Struct("<8sIII")
MAXIMUM_CASES = 64
OUTPUT_FLOAT_COUNT = 45

OUTPUT_NAMES = (
    "improvement_logit",
    "improvement_probability",
    "expected_defect_gain",
    "improvement_probability_variance",
    "expected_defect_gain_variance",
    *(f"plane_probability[{index}]" for index in range(training.FACE_COUNT)),
    *(f"plane_probability_variance[{index}]" for index in range(training.FACE_COUNT)),
    *(f"move_probability[{index}]" for index in range(training.MOVE_COUNT)),
    *(f"move_probability_variance[{index}]" for index in range(training.MOVE_COUNT)),
    *(f"scale_probability[{index}]" for index in range(training.SCALE_COUNT)),
    *(f"scale_probability_variance[{index}]" for index in range(training.SCALE_COUNT)),
)


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    repository = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(
        description=(
            "Compare a serialized SZTRNK01 model through CPU FP32 PyTorch and "
            "the native TransformerRanker implementation."
        )
    )
    parser.add_argument("--repository", default=str(repository))
    parser.add_argument(
        "--model", default="results/search/neural/transformer/current.sztf"
    )
    parser.add_argument(
        "--native-probe",
        help="TransformerRankerSelfTest executable containing --parity-probe",
    )
    parser.add_argument("--cases", type=int, default=8)
    parser.add_argument("--absolute-tolerance", type=float, default=2.0e-4)
    parser.add_argument("--relative-tolerance", type=float, default=2.0e-4)
    parser.add_argument("--timeout-seconds", type=float, default=30.0)
    parser.add_argument("--torch-threads", type=int, default=1)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    return parser.parse_args(argv)


def resolve_path(repository: Path, value: str) -> Path:
    path = Path(value).expanduser()
    if not path.is_absolute():
        path = repository / path
    return path.resolve()


def find_native_probe(repository: Path, explicit: str | None) -> Path:
    if explicit:
        path = resolve_path(repository, explicit)
        if not path.is_file():
            raise FileNotFoundError(f"native probe does not exist: {path}")
        return path
    candidates = (
        repository
        / "build/msbuild/bin/x64/Release/TransformerRankerSelfTest.exe",
        repository / "build/vs2026/Release/transformer_ranker_selftest.exe",
        repository / "build/TransformerRankerSelfTest.exe",
    )
    for candidate in candidates:
        if candidate.is_file():
            return candidate.resolve()
    raise FileNotFoundError(
        "native probe was not found; build transformer_ranker_selftest or pass "
        "--native-probe"
    )


def validate_options(args: argparse.Namespace) -> None:
    if not 1 <= args.cases <= MAXIMUM_CASES:
        raise ValueError(f"--cases must be in 1..{MAXIMUM_CASES}")
    for name in ("absolute_tolerance", "relative_tolerance"):
        value = float(getattr(args, name))
        if not math.isfinite(value) or value < 0.0:
            raise ValueError(f"--{name.replace('_', '-')} must be finite and non-negative")
    if not math.isfinite(args.timeout_seconds) or not 1.0 <= args.timeout_seconds <= 300.0:
        raise ValueError("--timeout-seconds must be in 1..300")
    if not 1 <= args.torch_threads <= 64:
        raise ValueError("--torch-threads must be in 1..64")


def load_model(
    path: Path, device: training.torch.device
) -> tuple[list[training.RankerMember], str]:
    model_bytes = path.read_bytes()
    expected_float_count = training.expected_payload_float_count()
    expected_payload_bytes = expected_float_count * 4
    expected_size = training.MODEL_HEADER_BYTES + expected_payload_bytes
    if len(model_bytes) != expected_size:
        raise ValueError(
            f"model size mismatch: {len(model_bytes)} != {expected_size}"
        )
    if model_bytes[:8] != training.MODEL_MAGIC:
        raise ValueError("model magic is invalid")

    format_version, feature_version, objective_version, approved = struct.unpack_from(
        "<4I", model_bytes, 8
    )
    dimensions = struct.unpack_from("<9I", model_bytes, 24)
    expected_dimensions = (
        training.FACE_COUNT,
        training.FACE_FEATURES,
        training.GLOBAL_FEATURES,
        training.MODEL_WIDTH,
        training.ATTENTION_HEADS,
        training.LAYER_COUNT,
        training.FEED_FORWARD_WIDTH,
        training.ENSEMBLE_SIZE,
        training.TOPOLOGY_COUNT,
    )
    if (
        format_version != training.MODEL_FORMAT_VERSION
        or feature_version != training.FEATURE_FORMAT_VERSION
        or objective_version != training.OBJECTIVE_VERSION
        or approved != 1
        or dimensions != expected_dimensions
        or struct.unpack_from("<I", model_bytes, 60)[0] != 0
        or not any(model_bytes[72:104])
        or struct.unpack_from("<QQ", model_bytes, 144) != (0, 0)
    ):
        raise ValueError("model header versions, dimensions, or approval are invalid")

    validation = struct.unpack_from("<4f", model_bytes, 104)
    if not all(math.isfinite(value) for value in validation):
        raise ValueError("model validation metadata is non-finite")
    if not all(0.0 <= validation[index] <= 1.0 for index in range(3)):
        raise ValueError("model validation probability metrics are out of range")

    payload_bytes, payload_crc, header_crc, float_count = struct.unpack_from(
        "<QIIQ", model_bytes, 120
    )
    if payload_bytes != expected_payload_bytes or float_count != expected_float_count:
        raise ValueError("model payload accounting is invalid")
    header = bytearray(model_bytes[: training.MODEL_HEADER_BYTES])
    struct.pack_into("<I", header, 132, 0)
    if zlib.crc32(header) & 0xFFFFFFFF != header_crc:
        raise ValueError("model header CRC-32 mismatch")
    payload = model_bytes[training.MODEL_HEADER_BYTES :]
    if zlib.crc32(payload) & 0xFFFFFFFF != payload_crc:
        raise ValueError("model payload CRC-32 mismatch")

    values = np.frombuffer(payload, dtype="<f4").copy()
    if values.size != expected_float_count or not np.isfinite(values).all():
        raise ValueError("model payload contains invalid float values")
    members: list[training.RankerMember] = []
    cursor = 0
    with training.torch.no_grad():
        for _ in range(training.ENSEMBLE_SIZE):
            member = training.RankerMember()
            for tensor in training.member_tensors(member):
                count = tensor.numel()
                source = training.torch.from_numpy(
                    values[cursor : cursor + count].reshape(tuple(tensor.shape))
                )
                tensor.copy_(source)
                cursor += count
            members.append(member.eval().to(device))
    if cursor != values.size:
        raise AssertionError(f"model tensor cursor mismatch: {cursor} != {values.size}")
    return members, hashlib.sha256(model_bytes).hexdigest()


def make_cases(case_count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    # Every value is an integer divided by 64, so both sides receive the same
    # exactly representable float32 bits. The bounded formula covers signs,
    # magnitudes, face positions and several topology embeddings.
    sample = np.arange(case_count, dtype=np.int64)[:, None, None]
    face = np.arange(training.FACE_COUNT, dtype=np.int64)[None, :, None]
    component = np.arange(training.FACE_FEATURES, dtype=np.int64)[None, None, :]
    face_features = (
        ((sample * 37 + face * 17 + component * 13) % 129) - 64
    ).astype(np.float32) / np.float32(64.0)
    global_component = np.arange(training.GLOBAL_FEATURES, dtype=np.int64)[None, :]
    global_features = (
        (
            (np.arange(case_count, dtype=np.int64)[:, None] * 29
             + global_component * 11)
            % 129
        )
        - 64
    ).astype(np.float32) / np.float32(64.0)
    topology = (
        (np.arange(case_count, dtype=np.uint32) * np.uint32(17) + np.uint32(42))
        % np.uint32(training.TOPOLOGY_COUNT)
    )
    return (
        np.ascontiguousarray(face_features, dtype="<f4"),
        np.ascontiguousarray(global_features, dtype="<f4"),
        np.ascontiguousarray(topology, dtype="<u4"),
    )


def write_probe_input(
    path: Path,
    face: np.ndarray,
    global_features: np.ndarray,
    topology: np.ndarray,
) -> None:
    case_count = int(face.shape[0])
    payload = bytearray(
        PROBE_INPUT_HEADER.pack(
            PROBE_INPUT_MAGIC,
            PROBE_VERSION,
            case_count,
            training.FACE_COUNT,
            training.FACE_FEATURES,
            training.GLOBAL_FEATURES,
        )
    )
    for index in range(case_count):
        payload.extend(struct.pack("<I", int(topology[index])))
        payload.extend(face[index].tobytes(order="C"))
        payload.extend(global_features[index].tobytes(order="C"))
    path.write_bytes(payload)


def read_probe_output(path: Path, expected_cases: int) -> tuple[np.ndarray, np.ndarray]:
    data = path.read_bytes()
    if len(data) < PROBE_OUTPUT_HEADER.size:
        raise ValueError("native probe output is truncated")
    magic, version, case_count, float_count = PROBE_OUTPUT_HEADER.unpack_from(data)
    if (
        magic != PROBE_OUTPUT_MAGIC
        or version != PROBE_VERSION
        or case_count != expected_cases
        or float_count != OUTPUT_FLOAT_COUNT
    ):
        raise ValueError("native probe output header is invalid")
    record_size = 4 + OUTPUT_FLOAT_COUNT * 4
    if len(data) != PROBE_OUTPUT_HEADER.size + case_count * record_size:
        raise ValueError("native probe output has trailing or truncated data")
    finite = np.empty(case_count, dtype=np.bool_)
    output = np.empty((case_count, OUTPUT_FLOAT_COUNT), dtype=np.float32)
    offset = PROBE_OUTPUT_HEADER.size
    for index in range(case_count):
        flag = struct.unpack_from("<I", data, offset)[0]
        if flag not in (0, 1):
            raise ValueError("native probe emitted an invalid finite flag")
        finite[index] = bool(flag)
        offset += 4
        output[index] = np.frombuffer(
            data, dtype="<f4", count=OUTPUT_FLOAT_COUNT, offset=offset
        )
        offset += OUTPUT_FLOAT_COUNT * 4
    if not np.isfinite(output).all():
        raise ValueError("native probe emitted NaN or infinity")
    return finite, output


@training.torch.inference_mode()
def pytorch_reference(
    members: list[training.RankerMember],
    face: np.ndarray,
    global_features: np.ndarray,
    topology: np.ndarray,
    device: training.torch.device,
) -> np.ndarray:
    face_tensor = training.torch.from_numpy(face).to(device)
    global_tensor = training.torch.from_numpy(global_features).to(device)
    topology_tensor = training.torch.from_numpy(topology.astype(np.int64)).to(device)
    probabilities = []
    gains = []
    plane_probabilities = []
    move_probabilities = []
    scale_probabilities = []
    for member in members:
        value, plane, move, scale = member(
            face_tensor, global_tensor, topology_tensor
        )
        probabilities.append(training.torch.sigmoid(value[:, 0]))
        gains.append(value[:, 1])
        plane_probabilities.append(training.torch.softmax(plane, dim=-1))
        move_probabilities.append(training.torch.softmax(move, dim=-1))
        scale_probabilities.append(training.torch.softmax(scale, dim=-1))

    def aggregate(values: list[training.torch.Tensor]):
        stacked = training.torch.stack(values)
        mean = stacked.mean(dim=0)
        variance = ((stacked - mean.unsqueeze(0)) ** 2).mean(dim=0)
        return mean, variance

    probability, probability_variance = aggregate(probabilities)
    gain, gain_variance = aggregate(gains)
    plane, plane_variance = aggregate(plane_probabilities)
    move, move_variance = aggregate(move_probabilities)
    scale, scale_variance = aggregate(scale_probabilities)
    improvement_logit = training.torch.logit(
        probability.clamp(1.0e-6, 1.0 - 1.0e-6)
    )
    result = training.torch.cat(
        (
            improvement_logit[:, None],
            probability[:, None],
            gain[:, None],
            probability_variance[:, None],
            gain_variance[:, None],
            plane,
            plane_variance,
            move,
            move_variance,
            scale,
            scale_variance,
        ),
        dim=1,
    )
    if result.shape[1] != OUTPUT_FLOAT_COUNT or not result.isfinite().all():
        raise RuntimeError("PyTorch reference emitted invalid output")
    return result.cpu().numpy()


def report_comparison(
    native: np.ndarray,
    reference: np.ndarray,
    absolute_tolerance: float,
    relative_tolerance: float,
) -> bool:
    difference = np.abs(native.astype(np.float64) - reference.astype(np.float64))
    tolerance = absolute_tolerance + relative_tolerance * np.abs(
        reference.astype(np.float64)
    )
    max_absolute_flat = int(np.argmax(difference))
    absolute_case, absolute_field = np.unravel_index(
        max_absolute_flat, difference.shape
    )
    relative = difference / np.maximum(
        np.maximum(np.abs(native), np.abs(reference)).astype(np.float64), 1.0e-12
    )
    max_relative_flat = int(np.argmax(relative))
    relative_case, relative_field = np.unravel_index(
        max_relative_flat, relative.shape
    )
    ratio = difference / np.maximum(tolerance, np.finfo(np.float64).tiny)
    max_ratio_flat = int(np.argmax(ratio))
    ratio_case, ratio_field = np.unravel_index(max_ratio_flat, ratio.shape)

    print(
        "max_abs_diff="
        f"{difference[absolute_case, absolute_field]:.9e} "
        f"case={absolute_case} field={OUTPUT_NAMES[absolute_field]} "
        f"native={native[absolute_case, absolute_field]:.9e} "
        f"pytorch={reference[absolute_case, absolute_field]:.9e}"
    )
    print(
        "max_rel_diff="
        f"{relative[relative_case, relative_field]:.9e} "
        f"case={relative_case} field={OUTPUT_NAMES[relative_field]} "
        f"native={native[relative_case, relative_field]:.9e} "
        f"pytorch={reference[relative_case, relative_field]:.9e}"
    )
    print(
        "max_tolerance_ratio="
        f"{ratio[ratio_case, ratio_field]:.9e} "
        f"case={ratio_case} field={OUTPUT_NAMES[ratio_field]}"
    )
    return bool(np.all(difference <= tolerance))


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_args(argv or sys.argv[1:])
    validate_options(args)
    repository = Path(args.repository).expanduser().resolve()
    model_path = resolve_path(repository, args.model)
    native_probe = find_native_probe(repository, args.native_probe)
    if not model_path.is_file():
        raise FileNotFoundError(f"model does not exist: {model_path}")
    if args.device == "cuda" and not training.torch.cuda.is_available():
        raise RuntimeError("--device cuda requested but CUDA PyTorch is unavailable")

    training.torch.set_num_threads(args.torch_threads)
    training.torch.set_float32_matmul_precision("highest")
    if training.torch.cuda.is_available():
        training.torch.backends.cuda.matmul.allow_tf32 = False
        training.torch.backends.cudnn.allow_tf32 = False
    device = training.torch.device(args.device)
    members, model_digest = load_model(model_path, device)
    face, global_features, topology = make_cases(args.cases)
    reference = pytorch_reference(
        members, face, global_features, topology, device
    )

    with tempfile.TemporaryDirectory(prefix="szilassi_transformer_parity_") as temp:
        input_path = Path(temp) / "input.bin"
        output_path = Path(temp) / "native.bin"
        write_probe_input(input_path, face, global_features, topology)
        command = [
            str(native_probe),
            "--parity-probe",
            str(model_path),
            str(input_path),
            str(output_path),
        ]
        completed = subprocess.run(
            command,
            cwd=repository,
            check=False,
            capture_output=True,
            text=True,
            timeout=args.timeout_seconds,
        )
        if completed.returncode != 0:
            raise RuntimeError(
                f"native probe failed with exit code {completed.returncode}: "
                f"{completed.stderr.strip() or completed.stdout.strip()}"
            )
        finite, native = read_probe_output(output_path, args.cases)
    if not finite.all():
        bad_cases = np.flatnonzero(~finite).tolist()
        raise RuntimeError(f"native inference fell back for cases: {bad_cases}")

    print(f"model_sha256={model_digest}")
    print(
        f"device={device.type} cases={args.cases} "
        f"compared_floats={native.size} torch_threads={args.torch_threads}"
    )
    passed = report_comparison(
        native,
        reference,
        args.absolute_tolerance,
        args.relative_tolerance,
    )
    print(
        f"tolerance=atol:{args.absolute_tolerance:.9e},"
        f"rtol:{args.relative_tolerance:.9e} status={'PASS' if passed else 'FAIL'}"
    )
    return 0 if passed else 1


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, ValueError, RuntimeError, subprocess.TimeoutExpired) as exc:
        print(f"Transformer parity check failed: {exc}", file=sys.stderr)
        raise SystemExit(2) from exc