"""
SoundBlockFabric — Unified API Tests

Tests all 15 primitives through the unified Fabric API.

Run:
    cd afs-engine
    python3 tests/test_sb_fabric_api.py -v
"""

import unittest
import numpy as np
import time
import sys
import os

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

from afs.gpu.sb_fabric import SoundBlockFabric, gpu_available


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricMap(unittest.TestCase):
    """Primitive 1: Carrier Modulation."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_abs(self):
        data = np.array([-3, -1, 0, 1, 3], dtype=np.float32)
        result = self.fabric.map(data, 'abs')
        np.testing.assert_array_almost_equal(result, [3, 1, 0, 1, 3])

    def test_tanh(self):
        data = np.array([-10, -1, 0, 1, 10], dtype=np.float32)
        result = self.fabric.map(data, 'tanh')
        expected = np.tanh(data)
        np.testing.assert_array_almost_equal(result, expected, decimal=3)

    def test_resonance(self):
        data = np.array([0, 1, 2, 3, 4], dtype=np.float32)
        result = self.fabric.map(data, 'resonance', scalar=2.0)
        expected = 1.0 / (1.0 + np.abs(data - 2.0))
        np.testing.assert_array_almost_equal(result, expected, decimal=3)

    def test_scale(self):
        data = np.array([1, 2, 3], dtype=np.float32)
        result = self.fabric.map(data, 'scale', scalar=3.0)
        np.testing.assert_array_almost_equal(result, [3, 6, 9])

    def test_square(self):
        data = np.array([1, 2, 3, 4], dtype=np.float32)
        result = self.fabric.map(data, 'square')
        np.testing.assert_array_almost_equal(result, [1, 4, 9, 16])


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricReduce(unittest.TestCase):
    """Primitive 2: Peak Detection."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_max(self):
        data = np.array([3, 1, 4, 1, 5, 9, 2, 6], dtype=np.float32)
        result = self.fabric.reduce(data, 'max')
        self.assertAlmostEqual(result, 9.0, places=2)

    def test_sum(self):
        data = np.ones(100, dtype=np.float32)
        result = self.fabric.reduce(data, 'sum')
        self.assertAlmostEqual(result, 100.0, places=0)

    def test_min(self):
        data = np.array([3, 1, 4, 1, 5, 9, 2, 6], dtype=np.float32)
        result = self.fabric.reduce(data, 'min')
        self.assertAlmostEqual(result, 1.0, places=2)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricScan(unittest.TestCase):
    """Primitive 3: Harmonic Accumulation."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_prefix_sum(self):
        data = np.array([1, 2, 3, 4], dtype=np.float32)
        result = self.fabric.scan(data, op='sum')
        np.testing.assert_array_almost_equal(result, [0, 1, 3, 6], decimal=2)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricSearch(unittest.TestCase):
    """Primitive 5: Resonance Matching."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_exact_match(self):
        data = np.array([10, 20, 30, 40, 50], dtype=np.float32)
        result = self.fabric.search(data, 30.0)
        self.assertEqual(result['index'], 2)
        self.assertAlmostEqual(result['score'], 1.0, places=2)

    def test_nearest(self):
        data = np.array([10, 20, 30, 40, 50], dtype=np.float32)
        result = self.fabric.search(data, 31.0)
        self.assertEqual(result['index'], 2)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricMatMul(unittest.TestCase):
    """Primitive 6: Carrier Matrix Modulation."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240, max_matrix_dim=256)

    def test_identity(self):
        A = np.eye(4, dtype=np.float32)
        B = np.array([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=np.float32)
        C = self.fabric.matmul(A, B)
        np.testing.assert_array_almost_equal(C, B, decimal=3)

    def test_square(self):
        A = np.random.randn(32, 32).astype(np.float32)
        B = np.random.randn(32, 32).astype(np.float32)
        C = self.fabric.matmul(A, B)
        expected = A @ B
        np.testing.assert_array_almost_equal(C, expected, decimal=1)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricFFT(unittest.TestCase):
    """Primitive 9: Spectral Transform."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_round_trip(self):
        signal = np.random.randn(64).astype(np.float32)
        re_freq, im_freq = self.fabric.fft(signal)
        re_back, _ = self.fabric.fft(re_freq, im_freq, inverse=True)
        np.testing.assert_array_almost_equal(re_back, signal, decimal=2)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricHistogram(unittest.TestCase):
    """Primitive 11: Spectral Energy Distribution."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_uniform(self):
        data = np.random.uniform(0, 1, 1000).astype(np.float32)
        counts = self.fabric.histogram(data, num_bins=10, val_range=(0, 1))
        self.assertEqual(int(np.sum(counts)), 1000)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricConvolve(unittest.TestCase):
    """Primitive 12: Resonance Filter."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_identity(self):
        data = np.array([1, 2, 3, 4, 5], dtype=np.float32)
        kernel = np.array([0, 0, 1, 0, 0], dtype=np.float32)
        result = self.fabric.convolve(data, kernel)
        np.testing.assert_array_almost_equal(result, data, decimal=3)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricCompact(unittest.TestCase):
    """Primitive 14: Active Bin Selection."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_remove_zeros(self):
        data = np.array([0, 1, 0, 2, 0, 3], dtype=np.float32)
        result = self.fabric.compact(data, threshold=0.5)
        np.testing.assert_array_almost_equal(result, [1, 2, 3])


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricSegmented(unittest.TestCase):
    """Primitive 13: Multi-Band Processing."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)

    def test_sum_per_segment(self):
        data = np.array([1, 2, 3, 4, 5, 6], dtype=np.float32)
        segments = np.array([0, 0, 1, 1, 2, 2], dtype=np.int64)
        result = self.fabric.segmented_reduce(data, segments, 'sum')
        np.testing.assert_array_almost_equal(result, [3, 7, 11])


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricGPUResident(unittest.TestCase):
    """GPU-resident slot operations — zero CPU transfer between ops."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=102400)

    def test_upload_download_roundtrip(self):
        """Upload → download recovers original data."""
        data = np.array([1, 2, 3, 4, 5], dtype=np.float32)
        self.fabric.upload('a', data)
        result = self.fabric.download('a')
        np.testing.assert_array_almost_equal(result, data)

    def test_map_gpu_tanh(self):
        """GPU-resident tanh — no CPU transfer."""
        data = np.array([-2, -1, 0, 1, 2], dtype=np.float32)
        self.fabric.upload('a', data)
        self.fabric.map_gpu('a', 'out', 'tanh')
        result = self.fabric.download('out')
        np.testing.assert_array_almost_equal(result, np.tanh(data), decimal=3)

    def test_chained_operations(self):
        """Chain: upload → tanh → square → download. GPU-resident throughout."""
        data = np.array([0.5, 1.0, 1.5, 2.0], dtype=np.float32)
        self.fabric.upload('a', data)
        self.fabric.map_gpu('a', 'out', 'tanh')       # out = tanh(a)
        self.fabric.map_gpu('out', 'a', 'square')      # a = out²
        result = self.fabric.download('a')
        expected = np.tanh(data) ** 2
        np.testing.assert_array_almost_equal(result, expected, decimal=3)

    def test_reduce_gpu(self):
        """GPU-resident reduce — only 4 bytes back to CPU."""
        data = np.array([3, 1, 4, 1, 5, 9, 2, 6], dtype=np.float32)
        self.fabric.upload('a', data)
        result = self.fabric.reduce_gpu('a', 'max')
        self.assertAlmostEqual(result, 9.0, places=2)

    def test_matmul_gpu(self):
        """GPU-resident matmul — zero transfer."""
        A = np.eye(4, dtype=np.float32)
        B = np.array([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=np.float32)
        self.fabric.upload('mat_a', A)
        self.fabric.upload('mat_b', B)
        self.fabric.matmul_gpu('mat_a', 'mat_b', 'mat_c', M=4, K=4, N=2)
        result = self.fabric.download('mat_c').reshape(4, 2)
        np.testing.assert_array_almost_equal(result, B, decimal=3)

    def test_large_chain_performance(self):
        """10K elements × 100 chained ops — should be <100ms total."""
        N = 10000
        data = np.random.randn(N).astype(np.float32)
        self.fabric.upload('a', data)

        t0 = time.perf_counter()
        for _ in range(100):
            self.fabric.map_gpu('a', 'out', 'tanh')
            self.fabric.map_gpu('out', 'a', 'scale', scalar=0.99)
        elapsed = (time.perf_counter() - t0) * 1000

        result = self.fabric.download('a')
        self.assertEqual(len(result), N)
        print(f"\n  GPU-resident chain: 100 ops × 10K elements = {elapsed:.1f}ms")
        print(f"  Per-op: {elapsed / 100:.2f}ms (vs 20ms with upload/download)")
        # Should be WAY faster than 100 × 20ms = 2000ms
        self.assertLess(elapsed, 500, "GPU-resident chain should be <500ms for 100 ops")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricPrimitivesList(unittest.TestCase):
    """Meta: verify all 20 primitives are listed (15 compute + 5 Virtual FPGA)."""

    def test_15_primitives(self):
        fabric = SoundBlockFabric(max_size=1024)
        self.assertEqual(len(fabric.primitives), 20)
        print(f"\n  Sound Block Fabric — {len(fabric.primitives)} primitives:")
        for p in fabric.primitives:
            print(f"    {p}")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricBenchmark(unittest.TestCase):
    """Performance benchmarks."""

    def test_benchmark_all(self):
        fabric = SoundBlockFabric(max_size=102400, max_matrix_dim=256)
        N = 10000
        data = np.random.randn(N).astype(np.float32)
        matrix = np.random.randn(64, 64).astype(np.float32)

        results = {}

        # Map
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.map(data, 'tanh')
        results['map (tanh, 10K)'] = (time.perf_counter() - t0) / 10 * 1000

        # Reduce
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.reduce(data, 'sum')
        results['reduce (sum, 10K)'] = (time.perf_counter() - t0) / 10 * 1000

        # Scan
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.scan(data[:1024], op='sum')
        results['scan (sum, 1K)'] = (time.perf_counter() - t0) / 10 * 1000

        # MatMul
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.matmul(matrix, matrix)
        results['matmul (64×64)'] = (time.perf_counter() - t0) / 10 * 1000

        # FFT
        fft_data = np.random.randn(1024).astype(np.float32)
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.fft(fft_data)
        results['fft (1024)'] = (time.perf_counter() - t0) / 10 * 1000

        # Histogram
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.histogram(data, num_bins=64, val_range=(-3, 3))
        results['histogram (10K, 64 bins)'] = (time.perf_counter() - t0) / 10 * 1000

        # Convolve
        kernel = np.ones(7, dtype=np.float32) / 7
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.convolve(data, kernel)
        results['convolve (10K, 7-tap)'] = (time.perf_counter() - t0) / 10 * 1000

        # Search
        t0 = time.perf_counter()
        for _ in range(10):
            fabric.search(data, 0.5)
        results['search (10K)'] = (time.perf_counter() - t0) / 10 * 1000

        print(f"\n  Sound Block Fabric Benchmark (GPU):")
        print(f"  {'Primitive':>30} {'Time':>10}")
        print(f"  {'─'*30} {'─'*10}")
        for name, ms in sorted(results.items(), key=lambda x: x[1]):
            print(f"  {name:>30} {ms:>8.2f}ms")
        print(f"\n  Total calls: {fabric.stats['calls']}")


if __name__ == "__main__":
    unittest.main(verbosity=2)
