"""Benchmark for _vals_equal optimization (Target 5)"""

import sys
import time

import numpy as np


# Original implementation
def vals_equal_original(v1, v2):
    """Original implementation"""
    if np is not None and (
        isinstance(v1, np.ndarray) or isinstance(v2, np.ndarray)
    ):
        return np.array_equal(v1, v2)
    if isinstance(v1, (list, tuple)):
        return (
            isinstance(v2, (list, tuple))
            and len(v1) == len(v2)
            and all(vals_equal_original(e1, e2) for e1, e2 in zip(v1, v2))
        )
    if isinstance(v1, dict):
        return (
            isinstance(v2, dict)
            and set(v1.keys()) == set(v2.keys())
            and all(vals_equal_original(v1[k], v2[k]) for k in v1)
        )
    return v1 == v2


# Optimized implementation with early exits
def vals_equal_optimized(v1, v2):
    """Optimized with shape check before element-wise comparison"""
    if np is not None and (
        isinstance(v1, np.ndarray) or isinstance(v2, np.ndarray)
    ):
        # Early exit for type mismatch
        if not (isinstance(v1, np.ndarray) and isinstance(v2, np.ndarray)):
            return False
        # Early exit for shape mismatch
        if v1.shape != v2.shape:
            return False
        # Now do element-wise comparison (np.array_equal handles dtype conversion)
        return np.array_equal(v1, v2)
    if isinstance(v1, (list, tuple)):
        return (
            isinstance(v2, (list, tuple))
            and len(v1) == len(v2)
            and all(vals_equal_optimized(e1, e2) for e1, e2 in zip(v1, v2))
        )
    if isinstance(v1, dict):
        return (
            isinstance(v2, dict)
            and set(v1.keys()) == set(v2.keys())
            and all(vals_equal_optimized(v1[k], v2[k]) for k in v1)
        )
    return v1 == v2


def benchmark(func, v1, v2, iterations=10000):
    """Benchmark a vals_equal implementation"""
    total_time = 0
    for _ in range(iterations):
        start = time.perf_counter()
        result = func(v1, v2)
        total_time += time.perf_counter() - start

    avg_time = total_time / iterations
    return avg_time, result


def main():
    print("Benchmarking _vals_equal implementations")
    print("=" * 60)

    # Test cases
    test_cases = [
        (
            "Equal arrays (same data)",
            np.array([1, 2, 3, 4, 5] * 100),
            np.array([1, 2, 3, 4, 5] * 100),
        ),
        (
            "Different shapes (early exit)",
            np.array([1, 2, 3, 4, 5] * 100),
            np.array([1, 2, 3, 4] * 100),
        ),
        (
            "Different dtypes (but same values)",
            np.array([1, 2, 3, 4, 5] * 100, dtype=np.int32),
            np.array([1, 2, 3, 4, 5] * 100, dtype=np.float64),
        ),
        (
            "Different values (late difference)",
            np.array([1, 2, 3, 4, 5] * 100),
            np.array([1, 2, 3, 4, 6] * 100),
        ),
        ("Large equal arrays", np.arange(10000), np.arange(10000)),
        ("Array vs scalar (type mismatch)", np.array([1, 2, 3]), [1, 2, 3]),
    ]

    for name, v1, v2 in test_cases:
        print(f"\nTest: {name}")

        # Benchmark original
        orig_time, orig_result = benchmark(vals_equal_original, v1, v2)

        # Benchmark optimized
        opt_time, opt_result = benchmark(vals_equal_optimized, v1, v2)

        # Verify results match
        if orig_result != opt_result:
            print("  ERROR: Results don't match!")
            print(f"    Original: {orig_result}")
            print(f"    Optimized: {opt_result}")
            sys.exit(1)

        # Report results
        speedup = orig_time / opt_time if opt_time > 0 else float("inf")
        print(f"  Result: {orig_result}")
        print(f"  Original:  {orig_time * 1e6:.2f} µs")
        print(f"  Optimized: {opt_time * 1e6:.2f} µs")
        print(f"  Speedup:   {speedup:.2f}x")

    print("\n" + "=" * 60)
    print("All tests passed!")


if __name__ == "__main__":
    main()
