"""Tensor comparison utilities for debugging MLX vs PyTorch divergence."""

from dataclasses import dataclass
from typing import Optional

import mlx.core as mx
import numpy as np


@dataclass
class TensorComparison:
    """Result of comparing two tensors."""

    name: str
    shape_a: tuple
    shape_b: tuple
    cosine_sim: float
    max_abs_diff: float
    mean_abs_diff: float
    passed: bool
    note: str = ""

    def __str__(self) -> str:
        status = "PASS" if self.passed else "FAIL"
        return (
            f"[{status}] {self.name}: "
            f"cosine={self.cosine_sim:.6f}, "
            f"max_diff={self.max_abs_diff:.2e}, "
            f"mean_diff={self.mean_abs_diff:.2e}"
            f"{f' ({self.note})' if self.note else ''}"
        )


def compare_tensors(
    a: mx.array,
    b: mx.array,
    name: str = "tensor",
    atol: float = 1e-5,
    cos_threshold: float = 0.999,
) -> TensorComparison:
    """Compare two tensors with multiple metrics.

    Args:
        a: First tensor (MLX).
        b: Second tensor (reference, can be mx.array or will be converted).
        name: Label for this comparison.
        atol: Absolute tolerance for max_abs_diff pass criterion.
        cos_threshold: Minimum cosine similarity to pass.

    Returns:
        TensorComparison with all metrics.
    """
    if a.shape != b.shape:
        return TensorComparison(
            name=name,
            shape_a=tuple(a.shape),
            shape_b=tuple(b.shape),
            cosine_sim=0.0,
            max_abs_diff=float("inf"),
            mean_abs_diff=float("inf"),
            passed=False,
            note=f"Shape mismatch: {a.shape} vs {b.shape}",
        )

    a_flat = a.reshape(-1).astype(mx.float32)
    b_flat = b.reshape(-1).astype(mx.float32)

    diff = mx.abs(a_flat - b_flat)
    max_abs = diff.max().item()
    mean_abs = diff.mean().item()

    # Cosine similarity
    dot = mx.sum(a_flat * b_flat).item()
    norm_a = mx.sqrt(mx.sum(a_flat * a_flat)).item()
    norm_b = mx.sqrt(mx.sum(b_flat * b_flat)).item()
    cos_sim = dot / (norm_a * norm_b + 1e-12)

    passed = cos_sim >= cos_threshold and max_abs <= atol

    return TensorComparison(
        name=name,
        shape_a=tuple(a.shape),
        shape_b=tuple(b.shape),
        cosine_sim=cos_sim,
        max_abs_diff=max_abs,
        mean_abs_diff=mean_abs,
        passed=passed,
    )


def compare_report(comparisons: list[TensorComparison]) -> str:
    """Generate a summary report from multiple comparisons."""
    lines = ["=" * 70, "Tensor Comparison Report", "=" * 70]
    failed = []
    for c in comparisons:
        lines.append(str(c))
        if not c.passed:
            failed.append(c)
    lines.append("-" * 70)
    lines.append(
        f"Total: {len(comparisons)}, "
        f"Passed: {len(comparisons) - len(failed)}, "
        f"Failed: {len(failed)}"
    )
    if failed:
        lines.append(f"First failure: {failed[0].name}")
    lines.append("=" * 70)
    return "\n".join(lines)
