"""
AFS GPU Force Kernel Tests — validates WGSL compute shaders against NumPy reference.

Tests run on GPU if available, skip gracefully if not.
"""

import numpy as np
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from afs.gpu import gpu_available

# Skip all tests if no GPU
HAS_GPU = gpu_available()


def run_tests():
    passed = 0
    failed = 0

    def test(name, fn):
        nonlocal passed, failed
        print(f"test_{name}...", end=" ")
        if not HAS_GPU:
            print("SKIP (no GPU)")
            return
        try:
            fn()
            print("PASS")
            passed += 1
        except Exception as e:
            print(f"FAIL ({e})")
            failed += 1

    def test_gpu_detection():
        from afs.gpu import WgpuForceKernel
        kernel = WgpuForceKernel(num_dims=5)
        info = kernel.adapter_info
        assert "device" in info
        assert "backend_type" in info or "vendor" in info

    def test_similarity_forces_parity():
        """GPU similarity forces match NumPy reference."""
        from afs.gpu import WgpuForceKernel
        np.random.seed(42)
        C, D = 64, 9
        velocities = np.random.uniform(-0.3, 0.3, (C, D)).astype(np.float32)
        matrix = np.random.uniform(-1.0, 1.0, (D, D)).astype(np.float32)
        np.fill_diagonal(matrix, 0.0)
        force_scale = 0.15
        tradeoff_amp = 1.5

        # NumPy reference
        amp = matrix.copy()
        amp[amp < 0] *= tradeoff_amp
        np.fill_diagonal(amp, 0.0)
        expected = (velocities @ amp) * force_scale

        # GPU
        kernel = WgpuForceKernel(num_dims=D, max_cells=C)
        kernel.upload_matrix(matrix)
        kernel.upload_concepts(np.zeros((1, D), dtype=np.float32))

        # Use compute with zero gravity/noise to isolate matrix forces
        positions = np.random.uniform(0.1, 0.9, (C, D)).astype(np.float32)
        kernel.compute(positions, velocities, {
            "force_scale": force_scale,
            "tradeoff_amplifier": tradeoff_amp,
            "gravity_strength": 0.0,  # disable gravity
            "gravity_range": 0.0,
            "dt": 16.67,
            "momentum": 0.0,  # zero momentum so velocity = force * dt_norm
            "max_speed": 100.0,
            "boundary_margin": 0.0,
            "boundary_strength": 0.0,
            "thermal_noise": 0.0,
            "tick_count": 0,
        })
        _, gpu_vel = kernel.readback()

        # velocity = 0 * momentum + (matrix_force + 0) * dt_norm
        # So gpu_vel = matrix_force * dt_norm = expected * dt_norm
        dt_norm = 16.67 / 1000.0
        gpu_forces = gpu_vel / dt_norm
        assert np.allclose(gpu_forces, expected, atol=1e-3), \
            f"Max error: {np.max(np.abs(gpu_forces - expected)):.2e}"

    def test_concept_gravity_parity():
        """GPU concept gravity matches NumPy reference."""
        from afs.gpu import WgpuForceKernel
        np.random.seed(42)
        C, D, K = 64, 9, 5
        cell_pos = np.random.uniform(0.1, 0.9, (C, D)).astype(np.float32)
        concept_pos = np.random.uniform(0.2, 0.8, (K, D)).astype(np.float32)

        strength, sharpness, grange, softening = 0.3, 2.0, 1.5, 0.05
        dt_norm = 16.67 / 1000.0

        # NumPy reference
        deltas = concept_pos[None,:,:] - cell_pos[:,None,:]
        distances = np.linalg.norm(deltas, axis=2)
        in_range = (distances < grange) & (distances > 0.001)
        grav_mag = np.where(in_range,
            strength * distances / (softening**2 + distances**2 * sharpness) * dt_norm, 0.0)
        safe = np.maximum(np.linalg.norm(deltas, axis=2, keepdims=True), 1e-6)
        expected = np.sum(grav_mag[:,:,None] * (deltas / safe), axis=1)

        # GPU — disable everything except gravity
        kernel = WgpuForceKernel(num_dims=D, max_cells=C)
        matrix = np.zeros((D, D), dtype=np.float32)
        kernel.upload_matrix(matrix)
        kernel.upload_concepts(concept_pos)

        kernel.compute(cell_pos, np.zeros_like(cell_pos), {
            "force_scale": 0.0,
            "tradeoff_amplifier": 1.0,
            "gravity_strength": strength,
            "gravity_sharpness": sharpness,
            "gravity_range": grange,
            "softening": softening,
            "dt": 16.67,
            "momentum": 0.0,
            "max_speed": 100.0,
            "boundary_margin": 0.0,
            "boundary_strength": 0.0,
            "thermal_noise": 0.0,
            "tick_count": 0,
        })
        _, gpu_vel = kernel.readback()
        gpu_forces = gpu_vel / dt_norm

        assert np.allclose(gpu_forces, expected, atol=1e-4), \
            f"Max error: {np.max(np.abs(gpu_forces - expected)):.2e}"

    def test_gpu_resident_mode():
        """GPU-resident tick() produces same results as compute() + readback()."""
        from afs.gpu import WgpuForceKernel
        np.random.seed(42)
        C, D, K = 32, 7, 4
        cell_pos = np.random.uniform(0.1, 0.9, (C, D)).astype(np.float32)
        cell_vel = np.random.uniform(-0.05, 0.05, (C, D)).astype(np.float32)
        concept_pos = np.random.uniform(0.2, 0.8, (K, D)).astype(np.float32)
        matrix = np.random.uniform(-1.0, 1.0, (D, D)).astype(np.float32)
        np.fill_diagonal(matrix, 0.0)

        config = {
            "force_scale": 0.15, "tradeoff_amplifier": 1.5,
            "gravity_strength": 0.3, "gravity_sharpness": 2.0,
            "gravity_range": 1.5, "softening": 0.05,
            "dt": 16.67, "momentum": 0.7, "max_speed": 0.3,
            "boundary_margin": 0.05, "boundary_strength": 0.5,
            "thermal_noise": 0.0, "tick_count": 0,
        }

        # Path 1: compute() (uploads each tick)
        k1 = WgpuForceKernel(num_dims=D, max_cells=C)
        k1.upload_matrix(matrix)
        k1.upload_concepts(concept_pos)
        k1.compute(cell_pos, cell_vel, config)
        p1, v1 = k1.readback()

        # Path 2: upload_cells() + tick() (GPU-resident)
        k2 = WgpuForceKernel(num_dims=D, max_cells=C)
        k2.upload_matrix(matrix)
        k2.upload_concepts(concept_pos)
        k2.upload_cells(cell_pos, cell_vel)
        k2.tick(config)
        p2, v2 = k2.readback()

        assert np.allclose(p1, p2, atol=1e-6), f"Position error: {np.max(np.abs(p1 - p2)):.2e}"
        assert np.allclose(v1, v2, atol=1e-6), f"Velocity error: {np.max(np.abs(v1 - v2)):.2e}"

    def test_multi_tick_stability():
        """100 GPU ticks produce no NaN or Inf values."""
        from afs.gpu import WgpuForceKernel
        np.random.seed(42)
        C, D, K = 128, 9, 5
        cell_pos = np.random.uniform(0.1, 0.9, (C, D)).astype(np.float32)
        cell_vel = np.zeros((C, D), dtype=np.float32)
        concept_pos = np.random.uniform(0.2, 0.8, (K, D)).astype(np.float32)
        matrix = np.random.uniform(-1.0, 1.0, (D, D)).astype(np.float32)
        np.fill_diagonal(matrix, 0.0)

        config = {
            "force_scale": 0.15, "tradeoff_amplifier": 1.5,
            "gravity_strength": 0.3, "gravity_sharpness": 2.0,
            "gravity_range": 1.5, "softening": 0.05,
            "dt": 16.67, "momentum": 0.7, "max_speed": 0.3,
            "boundary_margin": 0.05, "boundary_strength": 0.5,
            "thermal_noise": 0.001, "tick_count": 0,
        }

        kernel = WgpuForceKernel(num_dims=D, max_cells=C)
        kernel.upload_matrix(matrix)
        kernel.upload_concepts(concept_pos)
        kernel.upload_cells(cell_pos, cell_vel)

        for i in range(100):
            config["tick_count"] = i
            kernel.tick(config)

        pos, vel = kernel.readback()
        assert not np.any(np.isnan(pos)), "NaN in positions"
        assert not np.any(np.isnan(vel)), "NaN in velocities"
        assert not np.any(np.isinf(pos)), "Inf in positions"
        assert not np.any(np.isinf(vel)), "Inf in velocities"
        assert np.all(pos >= 0.0) and np.all(pos <= 1.0), "Positions out of [0,1]"

    def test_engine_gpu_cpu_parity():
        """AfsEngine with GPU produces similar results to CPU-only."""
        from afs.core import FieldDefinition, AfsEngine, AfsSimulationConfig

        np.random.seed(42)
        D = 5
        field = FieldDefinition()
        for i in range(D):
            did = field.add_dimension(name=f"dim_{i}", description=f"test dim {i}")
            field.promote_dimension(did)

        matrix = np.random.uniform(-1.0, 1.0, (D, D))
        np.fill_diagonal(matrix, 0.0)
        field.update_similarity_matrix(matrix)

        # GPU engine
        cfg1 = AfsSimulationConfig()
        np.random.seed(42)
        e1 = AfsEngine(field, cfg1, use_gpu=True)
        e1.initialize()
        for _ in range(50):
            e1.tick()
        gpu_positions = e1.get_cell_positions()

        # CPU engine
        np.random.seed(42)
        e2 = AfsEngine(field, AfsSimulationConfig(), use_gpu=False)
        e2.initialize()
        for _ in range(50):
            e2.tick()
        cpu_positions = e2.get_cell_positions()

        # Not exact match (different noise, f32 vs f64) but should have similar spread
        gpu_spread = np.std(gpu_positions)
        cpu_spread = np.std(cpu_positions)
        assert gpu_spread > 0.01, f"GPU cells didn't spread: std={gpu_spread}"
        assert cpu_spread > 0.01, f"CPU cells didn't spread: std={cpu_spread}"
        assert len(e1.cells) > 0, "GPU engine has no cells"
        assert len(e2.cells) > 0, "CPU engine has no cells"

    # Run all tests
    print("=" * 60)
    print("AFS GPU Force Kernel Tests")
    print("=" * 60)
    if not HAS_GPU:
        print("WARNING: No GPU detected — all tests will be skipped")
    print()

    test("gpu_detection", test_gpu_detection)
    test("similarity_forces_parity", test_similarity_forces_parity)
    test("concept_gravity_parity", test_concept_gravity_parity)
    test("gpu_resident_mode", test_gpu_resident_mode)
    test("multi_tick_stability", test_multi_tick_stability)
    test("engine_gpu_cpu_parity", test_engine_gpu_cpu_parity)

    print(f"\nResults: {passed} passed, {failed} failed")
    return failed == 0


if __name__ == "__main__":
    success = run_tests()
    sys.exit(0 if success else 1)
