"""Numerical parity validation utilities for cross-framework reliability."""

from __future__ import annotations

import hashlib
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Tuple

import numpy as np


DEFAULT_TOLERANCE_POLICY: Dict[str, Dict[str, float]] = {
    "float32": {"rtol": 1e-5, "atol": 1e-6},
    "float64": {"rtol": 1e-5, "atol": 1e-6},
    "float16": {"rtol": 5e-3, "atol": 5e-3},
    "bfloat16": {"rtol": 5e-3, "atol": 5e-3},
}


@dataclass
class ParityResult:
    """Structured parity output for framework compatibility checks."""

    passed: bool
    max_error: float
    failing_tensors: List[str] = field(default_factory=list)
    dtype_mismatches: List[str] = field(default_factory=list)
    checksum_mismatches: List[str] = field(default_factory=list)
    checked_tensors: int = 0

    def to_dict(self) -> Dict[str, object]:
        return {
            "passed": self.passed,
            "max_error": self.max_error,
            "failing_tensors": list(self.failing_tensors),
            "dtype_mismatches": list(self.dtype_mismatches),
            "checksum_mismatches": list(self.checksum_mismatches),
            "checked_tensors": self.checked_tensors,
        }


def _checksum(arr: np.ndarray) -> str:
    order = "F" if arr.flags.f_contiguous and not arr.flags.c_contiguous else "C"
    raw = arr.tobytes(order=order)
    return hashlib.sha256(raw).hexdigest()


def _tensor_tolerance(dtype_name: str) -> Tuple[float, float]:
    tol = DEFAULT_TOLERANCE_POLICY.get(dtype_name, None)
    if tol is None:
        return (0.0, 0.0)
    return (float(tol["rtol"]), float(tol["atol"]))


def compare_tensor_maps(
    reference: Mapping[str, np.ndarray],
    candidate: Mapping[str, np.ndarray],
) -> ParityResult:
    """Compare two tensor dictionaries under fixed dtype-aware tolerances."""
    failing_tensors: List[str] = []
    dtype_mismatches: List[str] = []
    checksum_mismatches: List[str] = []
    max_error = 0.0

    common = sorted(set(reference.keys()) & set(candidate.keys()))
    for name in common:
        ref = np.asarray(reference[name])
        out = np.asarray(candidate[name])
        if ref.shape != out.shape:
            failing_tensors.append(name)
            continue

        if str(ref.dtype) != str(out.dtype):
            dtype_mismatches.append(name)

        if np.issubdtype(ref.dtype, np.integer) or np.issubdtype(ref.dtype, np.bool_):
            if not np.array_equal(ref, out):
                failing_tensors.append(name)
                checksum_mismatches.append(name)
                continue
            if _checksum(ref) != _checksum(out):
                checksum_mismatches.append(name)
            continue

        rtol, atol = _tensor_tolerance(str(ref.dtype))
        abs_err = float(np.max(np.abs(ref.astype(np.float64) - out.astype(np.float64))))
        max_error = max(max_error, abs_err)
        if not np.allclose(ref, out, rtol=rtol, atol=atol):
            failing_tensors.append(name)
            continue
        if _checksum(ref) != _checksum(out):
            checksum_mismatches.append(name)

    missing = sorted(set(reference.keys()) - set(candidate.keys()))
    unexpected = sorted(set(candidate.keys()) - set(reference.keys()))
    failing_tensors.extend([f"missing:{name}" for name in missing])
    failing_tensors.extend([f"unexpected:{name}" for name in unexpected])

    return ParityResult(
        passed=len(failing_tensors) == 0,
        max_error=max_error,
        failing_tensors=failing_tensors,
        dtype_mismatches=dtype_mismatches,
        checksum_mismatches=checksum_mismatches,
        checked_tensors=len(common),
    )

