#!/usr/bin/env python3
"""
Verify CDNA Reader Against GGUF Direct

Compares tensors loaded from CDNA shards against GGUF-direct loading.
Must pass before wiring CDNA into generate.py.

Usage:
    python3 -m helix_cdc.regrow.verify_cdna_reader \
        --gguf ~/helix-cdc/models/gguf/tinyllama-1.1b-chat-v1.0.Q8_0.gguf \
        --cdna-dir /tmp/tinyllama_test \
        --tensor blk.0.attn_q.weight

    # Or verify all critical tensors:
    python3 -m helix_cdc.regrow.verify_cdna_reader \
        --gguf ~/helix-cdc/models/gguf/tinyllama-1.1b-chat-v1.0.Q8_0.gguf \
        --cdna-dir /tmp/tinyllama_test \
        --all-critical
"""

import argparse
import numpy as np
from pathlib import Path
import sys

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from helix_cdc.generate import load_weights_from_gguf
from helix_cdc.regrow.cdna_reader import CDNAModel


# Critical tensors to verify first
CRITICAL_TENSORS = [
    "token_embd.weight",
    "blk.0.attn_q.weight",
    "blk.0.attn_k.weight",
    "blk.0.attn_v.weight",
    "blk.0.attn_output.weight",
    "blk.0.ffn_gate.weight",
    "blk.0.ffn_up.weight",
    "blk.0.ffn_down.weight",
    "output.weight",
]


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Compute cosine similarity between two arrays."""
    a_flat = a.flatten().astype(np.float64)
    b_flat = b.flatten().astype(np.float64)

    dot = np.dot(a_flat, b_flat)
    norm_a = np.linalg.norm(a_flat)
    norm_b = np.linalg.norm(b_flat)

    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


def verify_tensor(
    tensor_name: str,
    gguf_weights: dict,
    cdna_model: CDNAModel,
) -> dict:
    """
    Verify a single tensor matches between GGUF and CDNA.

    Returns dict with:
        - tensor_name
        - gguf_shape
        - cdna_shape
        - max_abs_diff
        - cosine
        - passed
    """
    result = {
        "tensor_name": tensor_name,
        "gguf_shape": None,
        "cdna_shape": None,
        "max_abs_diff": None,
        "cosine": None,
        "passed": False,
        "error": None,
    }

    # Load GGUF tensor
    if tensor_name not in gguf_weights:
        result["error"] = f"Not found in GGUF weights"
        return result

    gguf_tensor = gguf_weights[tensor_name]
    result["gguf_shape"] = gguf_tensor.shape

    # Load CDNA tensor
    if not cdna_model.has(tensor_name):
        result["error"] = f"Not found in CDNA model"
        return result

    try:
        cdna_tensor = cdna_model.load_tensor(tensor_name)
    except Exception as e:
        result["error"] = f"Failed to load from CDNA: {e}"
        return result

    result["cdna_shape"] = cdna_tensor.shape

    # Check shape match
    if gguf_tensor.shape != cdna_tensor.shape:
        result["error"] = f"Shape mismatch: GGUF {gguf_tensor.shape} vs CDNA {cdna_tensor.shape}"
        return result

    # Compute metrics
    diff = np.abs(gguf_tensor - cdna_tensor)
    result["max_abs_diff"] = float(diff.max())
    result["cosine"] = cosine_similarity(gguf_tensor, cdna_tensor)

    # Pass criteria for k-means quantization:
    # - cosine > 0.99 (acceptable for 256-cluster k-means)
    # - Small tensors (like attn_k) can have slightly lower cosine
    # - max_diff can be larger for k-means (up to 3.0 is acceptable)
    if result["cosine"] > 0.99:
        result["passed"] = True
    elif result["cosine"] > 0.98 and result["max_abs_diff"] < 3.0:
        # Small tensors with high variance
        result["passed"] = True

    return result


def print_result(result: dict) -> None:
    """Print verification result in a clean format."""
    name = result["tensor_name"]
    status = "PASS" if result["passed"] else "FAIL"

    if result["error"]:
        print(f"[{status}] {name}")
        print(f"       Error: {result['error']}")
        return

    print(f"[{status}] {name}")
    print(f"       GGUF shape:   {result['gguf_shape']}")
    print(f"       CDNA shape:   {result['cdna_shape']}")
    print(f"       Max abs diff: {result['max_abs_diff']:.6f}")
    print(f"       Cosine:       {result['cosine']:.6f}")


def main():
    parser = argparse.ArgumentParser(
        description="Verify CDNA tensors match GGUF-direct tensors"
    )
    parser.add_argument(
        "--gguf",
        required=True,
        help="Path to GGUF model file",
    )
    parser.add_argument(
        "--cdna-dir",
        required=True,
        help="Path to CDNA shard directory",
    )
    parser.add_argument(
        "--tensor",
        help="Single tensor name to verify (e.g., blk.0.attn_q.weight)",
    )
    parser.add_argument(
        "--all-critical",
        action="store_true",
        help="Verify all critical tensors",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="Verify all available tensors",
    )

    args = parser.parse_args()

    gguf_path = Path(args.gguf).expanduser()
    cdna_dir = Path(args.cdna_dir).expanduser()

    # Load GGUF weights
    print(f"Loading GGUF: {gguf_path}")
    gguf_weights = load_weights_from_gguf(gguf_path)
    print(f"  Loaded {len(gguf_weights)} tensors")

    # Load CDNA model
    print(f"Loading CDNA: {cdna_dir}")
    cdna_model = CDNAModel(str(cdna_dir))
    print(f"  Found {len(cdna_model.tensor_names())} tensors")

    # Determine which tensors to verify
    if args.tensor:
        tensors_to_verify = [args.tensor]
    elif args.all_critical:
        tensors_to_verify = CRITICAL_TENSORS
    elif args.all:
        tensors_to_verify = cdna_model.tensor_names()
    else:
        print("Error: Must specify --tensor, --all-critical, or --all")
        return 1

    print()
    print("=" * 60)
    print("Verification Results")
    print("=" * 60)

    passed = 0
    failed = 0

    for tensor_name in tensors_to_verify:
        result = verify_tensor(tensor_name, gguf_weights, cdna_model)
        print_result(result)
        print()

        if result["passed"]:
            passed += 1
        else:
            failed += 1

    print("=" * 60)
    print(f"Summary: {passed} passed, {failed} failed")
    print("=" * 60)

    return 0 if failed == 0 else 1


if __name__ == "__main__":
    exit(main())
