#!/usr/bin/env python3
"""
Sanity test for helix64 DOT product numerical equivalence.

Usage:
    pytest tests/test_helix64_dot.py
    python3 tests/test_helix64_dot.py  # self-test mode
"""
import sys
import struct
from pathlib import Path

import numpy as np
import pytest

# Add scripts/ to path
REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "scripts"))

from helix64_exec import execute_real, OP


def test_helix64_dot_basic():
    """Test basic DOT product: r[30] = dot(r[10:26], r[30:46])"""
    # Build tiny DOT program (use non-overlapping registers)
    # DOT r50, r10, r30, len=16; HALT
    prog = bytes([OP["DOT"], 50, 10, 30]) + struct.pack("<H", 16) + bytes([OP["HALT"], 0, 0, 0])

    # Test vectors
    a = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
                  9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0], dtype=np.float64)
    b = np.array([16.0, 15.0, 14.0, 13.0, 12.0, 11.0, 10.0, 9.0,
                  8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0], dtype=np.float64)

    # Expected result (NumPy reference)
    ref_result = float(np.dot(a, b))

    # Seed helix64 registers (non-overlapping: r10-25 and r30-45)
    R = np.zeros(256, dtype=np.float64)
    R[10:10+16] = a
    R[30:30+16] = b

    # Execute
    result_dict = execute_real(prog, fuel=10000, init_registers=R)
    h64_result = float(result_dict["registers"][50])  # Read from destination register r50

    # Check equivalence
    diff = abs(ref_result - h64_result)
    eps = 1e-6

    print(f"Reference (NumPy): {ref_result}")
    print(f"helix64 result:    {h64_result}")
    print(f"Absolute diff:     {diff:.2e}")
    print(f"Threshold:         {eps:.2e}")

    assert diff < eps, f"DOT product failed: {diff:.2e} > {eps:.2e}"
    print("✅ DOT product test PASSED")


def test_helix64_dot_zeros():
    """Test DOT with zero vectors"""
    prog = bytes([OP["DOT"], 50, 10, 30]) + struct.pack("<H", 8) + bytes([OP["HALT"], 0, 0, 0])

    a = np.zeros(8, dtype=np.float64)
    b = np.zeros(8, dtype=np.float64)

    R = np.zeros(256, dtype=np.float64)
    R[10:10+8] = a
    R[30:30+8] = b

    result_dict = execute_real(prog, fuel=10000, init_registers=R)
    h64_result = float(result_dict["registers"][50])

    assert abs(h64_result) < 1e-12, f"Zero DOT failed: {h64_result}"
    print("✅ Zero DOT test PASSED")


def test_helix64_dot_negative():
    """Test DOT with negative values"""
    prog = bytes([OP["DOT"], 50, 10, 30]) + struct.pack("<H", 4) + bytes([OP["HALT"], 0, 0, 0])

    a = np.array([1.0, -2.0, 3.0, -4.0], dtype=np.float64)
    b = np.array([-1.0, 2.0, -3.0, 4.0], dtype=np.float64)

    ref_result = float(np.dot(a, b))

    R = np.zeros(256, dtype=np.float64)
    R[10:10+4] = a
    R[30:30+4] = b

    result_dict = execute_real(prog, fuel=10000, init_registers=R)
    h64_result = float(result_dict["registers"][50])

    diff = abs(ref_result - h64_result)
    assert diff < 1e-6, f"Negative DOT failed: {diff:.2e}"
    print("✅ Negative DOT test PASSED")


def test_helix64_dot_metrics():
    """Test that DOT produces metrics (cycles, flops, bytes)"""
    prog = bytes([OP["DOT"], 50, 10, 30]) + struct.pack("<H", 16) + bytes([OP["HALT"], 0, 0, 0])

    a = np.random.randn(16).astype(np.float64)
    b = np.random.randn(16).astype(np.float64)

    R = np.zeros(256, dtype=np.float64)
    R[10:10+16] = a
    R[30:30+16] = b

    result_dict = execute_real(prog, fuel=10000, init_registers=R)
    metrics = result_dict["metrics"]

    # Check metrics exist and are reasonable
    assert metrics["cycles"] > 0, "cycles should be > 0"
    assert metrics["flops"] == 16, f"Expected 16 FLOPs, got {metrics['flops']}"
    assert metrics["bytes"] == 256, f"Expected 256 bytes, got {metrics['bytes']}"

    print(f"Metrics: cycles={metrics['cycles']}, flops={metrics['flops']}, bytes={metrics['bytes']}")
    print("✅ Metrics test PASSED")


if __name__ == "__main__":
    # Self-test mode (no pytest required)
    print("Running helix64 DOT sanity tests...")
    print("=" * 60)

    try:
        test_helix64_dot_basic()
        print()
        test_helix64_dot_zeros()
        print()
        test_helix64_dot_negative()
        print()
        test_helix64_dot_metrics()
        print()
        print("=" * 60)
        print("✅ All tests PASSED")
        sys.exit(0)
    except AssertionError as e:
        print(f"❌ Test FAILED: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
