#!/usr/bin/env python3
"""
comparator.py - GGUF File Comparison Utilities
==============================================

Compares two GGUF files for tensor-level equivalence.
Useful for verifying CDNA → GGUF reconstruction fidelity.

Features:
- Tensor-by-tensor comparison
- Configurable tolerance (absolute, relative, cosine)
- Summary statistics
- Mismatch reporting
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

try:
    from helix_cdc.gguf.tensor_index import GGUFTensorIndex
    HAVE_INDEX = True
except ImportError:
    HAVE_INDEX = False


@dataclass
class TensorComparison:
    """Result of comparing a single tensor."""
    name: str
    shape_a: List[int]
    shape_b: List[int]
    dtype_a: str
    dtype_b: str

    # Metrics
    max_abs_diff: float = 0.0
    mean_abs_diff: float = 0.0
    cosine_similarity: float = 1.0
    relative_error: float = 0.0

    # Status
    shapes_match: bool = True
    passes_threshold: bool = True
    error: Optional[str] = None


@dataclass
class GGUFComparisonResult:
    """Result of comparing two GGUF files."""
    file_a: str
    file_b: str
    timestamp: str

    # Summary
    tensors_compared: int = 0
    tensors_matched: int = 0
    tensors_mismatched: int = 0
    tensors_missing_a: int = 0
    tensors_missing_b: int = 0

    # Thresholds used
    abs_threshold: float = 0.001
    cosine_threshold: float = 0.99

    # Details
    comparisons: List[TensorComparison] = field(default_factory=list)
    mismatches: List[str] = field(default_factory=list)
    missing_in_a: List[str] = field(default_factory=list)
    missing_in_b: List[str] = field(default_factory=list)

    @property
    def all_pass(self) -> bool:
        return (
            self.tensors_mismatched == 0
            and self.tensors_missing_a == 0
            and self.tensors_missing_b == 0
        )


class GGUFComparator:
    """
    Compare two GGUF files for tensor equivalence.

    Usage:
        comparator = GGUFComparator(
            abs_threshold=0.001,
            cosine_threshold=0.99,
        )

        result = comparator.compare("baseline.gguf", "rebuilt.gguf")

        if result.all_pass:
            print("Files are equivalent!")
        else:
            print(f"Mismatches: {result.mismatches}")

        # Save report
        comparator.save_report(result, "comparison_report.json")
    """

    def __init__(
        self,
        abs_threshold: float = 0.001,
        cosine_threshold: float = 0.99,
        verbose: bool = False,
    ):
        """
        Initialize comparator.

        Args:
            abs_threshold: Maximum absolute difference to pass
            cosine_threshold: Minimum cosine similarity to pass
            verbose: Print progress during comparison
        """
        if not HAVE_INDEX:
            raise ImportError("GGUFTensorIndex required. Install: pip install gguf")

        self.abs_threshold = abs_threshold
        self.cosine_threshold = cosine_threshold
        self.verbose = verbose

    def compare(
        self,
        file_a: Path,
        file_b: Path,
        tensor_filter: Optional[str] = None,
    ) -> GGUFComparisonResult:
        """
        Compare two GGUF files.

        Args:
            file_a: First GGUF file (baseline/reference)
            file_b: Second GGUF file (rebuilt/test)
            tensor_filter: Optional regex pattern to filter tensors

        Returns:
            GGUFComparisonResult with comparison details
        """
        file_a = Path(file_a)
        file_b = Path(file_b)

        # Build indices
        if self.verbose:
            print(f"Building index for {file_a.name}...")
        index_a = GGUFTensorIndex.from_gguf(file_a)

        if self.verbose:
            print(f"Building index for {file_b.name}...")
        index_b = GGUFTensorIndex.from_gguf(file_b)

        # Get tensor lists
        tensors_a = set(index_a.list_tensors(pattern=tensor_filter))
        tensors_b = set(index_b.list_tensors(pattern=tensor_filter))

        all_tensors = tensors_a | tensors_b

        result = GGUFComparisonResult(
            file_a=str(file_a),
            file_b=str(file_b),
            timestamp=datetime.now(timezone.utc).isoformat(),
            abs_threshold=self.abs_threshold,
            cosine_threshold=self.cosine_threshold,
        )

        # Track missing tensors
        result.missing_in_a = sorted(tensors_b - tensors_a)
        result.missing_in_b = sorted(tensors_a - tensors_b)
        result.tensors_missing_a = len(result.missing_in_a)
        result.tensors_missing_b = len(result.missing_in_b)

        # Compare common tensors
        common_tensors = sorted(tensors_a & tensors_b)

        for i, name in enumerate(common_tensors):
            if self.verbose:
                print(f"  [{i+1}/{len(common_tensors)}] Comparing {name}...")

            comparison = self._compare_tensor(index_a, index_b, name)
            result.comparisons.append(comparison)
            result.tensors_compared += 1

            if comparison.passes_threshold:
                result.tensors_matched += 1
            else:
                result.tensors_mismatched += 1
                result.mismatches.append(name)

        return result

    def _compare_tensor(
        self,
        index_a: GGUFTensorIndex,
        index_b: GGUFTensorIndex,
        name: str,
    ) -> TensorComparison:
        """Compare a single tensor between two files."""
        info_a = index_a.get_tensor(name)
        info_b = index_b.get_tensor(name)

        comparison = TensorComparison(
            name=name,
            shape_a=info_a["shape"],
            shape_b=info_b["shape"],
            dtype_a=info_a.get("ggml_type", "unknown"),
            dtype_b=info_b.get("ggml_type", "unknown"),
        )

        # Check shapes match
        if info_a["shape"] != info_b["shape"]:
            comparison.shapes_match = False
            comparison.passes_threshold = False
            comparison.error = f"Shape mismatch: {info_a['shape']} vs {info_b['shape']}"
            return comparison

        # Load tensors
        try:
            tensor_a = index_a.load_tensor(name)
            tensor_b = index_b.load_tensor(name)
        except Exception as e:
            comparison.error = f"Load error: {e}"
            comparison.passes_threshold = False
            return comparison

        # Compute metrics
        diff = np.abs(tensor_a - tensor_b)
        comparison.max_abs_diff = float(np.max(diff))
        comparison.mean_abs_diff = float(np.mean(diff))

        # Cosine similarity
        flat_a = tensor_a.flatten()
        flat_b = tensor_b.flatten()
        norm_a = np.linalg.norm(flat_a)
        norm_b = np.linalg.norm(flat_b)

        if norm_a > 0 and norm_b > 0:
            comparison.cosine_similarity = float(
                np.dot(flat_a, flat_b) / (norm_a * norm_b)
            )
        else:
            comparison.cosine_similarity = 1.0 if norm_a == norm_b else 0.0

        # Relative error
        if norm_a > 0:
            comparison.relative_error = float(np.linalg.norm(diff) / norm_a)
        else:
            comparison.relative_error = 0.0 if norm_b == 0 else float("inf")

        # Check thresholds
        comparison.passes_threshold = (
            comparison.max_abs_diff <= self.abs_threshold
            and comparison.cosine_similarity >= self.cosine_threshold
        )

        return comparison

    def save_report(
        self,
        result: GGUFComparisonResult,
        output_path: Path,
    ) -> None:
        """Save comparison report as JSON."""
        output_path = Path(output_path)

        report = {
            "schema": "gguf_comparison_v1",
            "file_a": result.file_a,
            "file_b": result.file_b,
            "timestamp": result.timestamp,
            "summary": {
                "all_pass": result.all_pass,
                "tensors_compared": result.tensors_compared,
                "tensors_matched": result.tensors_matched,
                "tensors_mismatched": result.tensors_mismatched,
                "tensors_missing_a": result.tensors_missing_a,
                "tensors_missing_b": result.tensors_missing_b,
            },
            "thresholds": {
                "abs_threshold": result.abs_threshold,
                "cosine_threshold": result.cosine_threshold,
            },
            "mismatches": result.mismatches,
            "missing_in_a": result.missing_in_a,
            "missing_in_b": result.missing_in_b,
            "comparisons": [
                {
                    "name": c.name,
                    "shapes_match": c.shapes_match,
                    "passes_threshold": c.passes_threshold,
                    "max_abs_diff": c.max_abs_diff,
                    "mean_abs_diff": c.mean_abs_diff,
                    "cosine_similarity": c.cosine_similarity,
                    "relative_error": c.relative_error,
                    "error": c.error,
                }
                for c in result.comparisons
            ],
        }

        with open(output_path, "w") as f:
            json.dump(report, f, indent=2)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Compare two GGUF files")
    parser.add_argument("file_a", help="First GGUF file (baseline)")
    parser.add_argument("file_b", help="Second GGUF file (test)")
    parser.add_argument(
        "--abs-threshold",
        type=float,
        default=0.001,
        help="Maximum absolute difference (default: 0.001)",
    )
    parser.add_argument(
        "--cosine-threshold",
        type=float,
        default=0.99,
        help="Minimum cosine similarity (default: 0.99)",
    )
    parser.add_argument(
        "--filter",
        type=str,
        help="Regex pattern to filter tensor names",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Output path for JSON report",
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Verbose output",
    )

    args = parser.parse_args()

    comparator = GGUFComparator(
        abs_threshold=args.abs_threshold,
        cosine_threshold=args.cosine_threshold,
        verbose=args.verbose,
    )

    result = comparator.compare(
        Path(args.file_a),
        Path(args.file_b),
        tensor_filter=args.filter,
    )

    print()
    print("=" * 60)
    print("GGUF COMPARISON RESULT")
    print("=" * 60)
    print(f"File A: {result.file_a}")
    print(f"File B: {result.file_b}")
    print()
    print(f"Tensors compared: {result.tensors_compared}")
    print(f"Tensors matched:  {result.tensors_matched}")
    print(f"Tensors mismatch: {result.tensors_mismatched}")
    print(f"Missing in A:     {result.tensors_missing_a}")
    print(f"Missing in B:     {result.tensors_missing_b}")
    print()

    if result.all_pass:
        print("RESULT: PASS - Files are equivalent")
    else:
        print("RESULT: FAIL")
        if result.mismatches:
            print(f"  Mismatched: {result.mismatches[:5]}...")
        if result.missing_in_a:
            print(f"  Missing in A: {result.missing_in_a[:5]}...")
        if result.missing_in_b:
            print(f"  Missing in B: {result.missing_in_b[:5]}...")

    if args.output:
        comparator.save_report(result, args.output)
        print(f"\nReport saved to: {args.output}")
