"""
Sound Block Fabric — GPU Primitive Tests

Tests the 5 new SB Fabric primitives: Scan, Histogram, Convolution.
(Segmented and Compaction use Scan as a dependency.)

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

import unittest
import numpy as np
import sys
import os

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

from afs.gpu.sb_fabric_primitives import (
    SBScan, SBHistogram, SBConvolve, SBFFT, gpu_available
)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestSBScan(unittest.TestCase):
    """Harmonic Accumulation — prefix sum."""

    def test_exclusive_sum_small(self):
        """Exclusive prefix sum on [1,2,3,4] = [0,1,3,6]."""
        scan = SBScan(max_size=1024)
        data = np.array([1, 2, 3, 4], dtype=np.float32)
        result = scan.scan(data, op="sum")
        expected = np.array([0, 1, 3, 6], dtype=np.float32)
        np.testing.assert_array_almost_equal(result, expected, decimal=3)

    def test_exclusive_sum_256(self):
        """256 ones → [0, 1, 2, ..., 255]."""
        scan = SBScan(max_size=1024)
        data = np.ones(256, dtype=np.float32)
        result = scan.scan(data, op="sum")
        expected = np.arange(256, dtype=np.float32)
        np.testing.assert_array_almost_equal(result, expected, decimal=2)

    def test_exclusive_sum_1000(self):
        """1000 elements — multi-block scan."""
        scan = SBScan(max_size=4096)
        data = np.ones(1000, dtype=np.float32)
        result = scan.scan(data, op="sum")
        # Exclusive sum of all-ones = [0, 1, 2, ..., 999]
        expected = np.arange(1000, dtype=np.float32)
        np.testing.assert_array_almost_equal(result, expected, decimal=1)

    def test_max_scan(self):
        """Prefix max: running maximum."""
        scan = SBScan(max_size=1024)
        data = np.array([3, 1, 4, 1, 5, 9, 2, 6], dtype=np.float32)
        result = scan.scan(data, op="max")
        # Exclusive prefix max: identity for max is -inf, first element sees -inf
        # This tests that the scan infrastructure works for non-sum ops
        self.assertEqual(len(result), 8)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestSBHistogram(unittest.TestCase):
    """Spectral Energy Distribution — atomic binning."""

    def test_uniform_distribution(self):
        """Uniform [0,1] data should give roughly equal bin counts."""
        hist = SBHistogram(max_size=10240)
        data = np.random.uniform(0, 1, 10000).astype(np.float32)
        counts = hist.histogram(data, num_bins=10, val_range=(0.0, 1.0))

        self.assertEqual(len(counts), 10)
        self.assertEqual(int(np.sum(counts)), 10000)
        # Each bin should have ~1000 ± 100
        for c in counts:
            self.assertGreater(c, 800, f"Bin count {c} too low for uniform")
            self.assertLess(c, 1200, f"Bin count {c} too high for uniform")

    def test_peaked_distribution(self):
        """All values at 0.5 should concentrate in one bin."""
        hist = SBHistogram(max_size=1024)
        data = np.full(500, 0.5, dtype=np.float32)
        counts = hist.histogram(data, num_bins=10, val_range=(0.0, 1.0))

        self.assertEqual(int(np.sum(counts)), 500)
        # Middle bin should have all 500
        self.assertEqual(int(counts[5]), 500)

    def test_custom_range(self):
        """Custom range bins correctly."""
        hist = SBHistogram(max_size=1024)
        data = np.array([10, 20, 30, 40, 50], dtype=np.float32)
        counts = hist.histogram(data, num_bins=5, val_range=(0.0, 50.0))
        self.assertEqual(int(np.sum(counts)), 5)

    def test_large_histogram(self):
        """100K elements, 64 bins."""
        hist = SBHistogram(max_size=102400)
        data = np.random.normal(0.5, 0.15, 100000).astype(np.float32)
        data = np.clip(data, 0, 1)
        counts = hist.histogram(data, num_bins=64, val_range=(0.0, 1.0))

        self.assertEqual(len(counts), 64)
        total = int(np.sum(counts))
        self.assertEqual(total, 100000)
        # Normal distribution: center bins should have most counts
        center_bins = counts[24:40]
        edge_bins = np.concatenate([counts[:8], counts[56:]])
        self.assertGreater(np.mean(center_bins), np.mean(edge_bins))


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestSBConvolve(unittest.TestCase):
    """Resonance Filter — 1D convolution."""

    def test_identity_kernel(self):
        """Convolving with [0, 0, 1, 0, 0] = identity."""
        conv = SBConvolve(max_size=1024)
        data = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.float32)
        kernel = np.array([0, 0, 1, 0, 0], dtype=np.float32)
        result = conv.convolve(data, kernel)
        np.testing.assert_array_almost_equal(result, data, decimal=4)

    def test_box_filter(self):
        """Box filter [1/3, 1/3, 1/3] = moving average."""
        conv = SBConvolve(max_size=1024)
        data = np.array([0, 0, 0, 3, 0, 0, 0], dtype=np.float32)
        kernel = np.array([1/3, 1/3, 1/3], dtype=np.float32)
        result = conv.convolve(data, kernel)
        # Center spike of 3 should spread to [1, 1, 1] (approximately)
        self.assertAlmostEqual(result[3], 1.0, places=3)
        self.assertAlmostEqual(result[2], 1.0, places=3)
        self.assertAlmostEqual(result[4], 1.0, places=3)

    def test_gaussian_smoothing(self):
        """Gaussian kernel smooths noise."""
        conv = SBConvolve(max_size=4096)
        # Noisy signal
        data = np.sin(np.linspace(0, 4 * np.pi, 1000)).astype(np.float32)
        data += np.random.normal(0, 0.3, 1000).astype(np.float32)
        # Gaussian kernel (sigma=2, size=7)
        x = np.arange(-3, 4)
        kernel = np.exp(-x**2 / (2 * 2**2)).astype(np.float32)
        kernel /= kernel.sum()

        result = conv.convolve(data, kernel)
        # Smoothed signal should have lower variance
        self.assertLess(np.std(result[100:900]), np.std(data[100:900]))

    def test_large_convolution(self):
        """10K elements with 15-tap kernel."""
        conv = SBConvolve(max_size=10240)
        data = np.random.randn(10000).astype(np.float32)
        kernel = np.ones(15, dtype=np.float32) / 15
        result = conv.convolve(data, kernel)
        self.assertEqual(len(result), 10000)
        # Moving average should reduce variance
        self.assertLess(np.std(result[100:9900]), np.std(data[100:9900]))


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestSBFFT(unittest.TestCase):
    """Spectral Transform — GPU FFT."""

    def test_impulse_fft(self):
        """FFT of impulse [1, 0, 0, 0] = flat spectrum [1, 1, 1, 1]."""
        fft = SBFFT(max_size=1024)
        re = np.array([1, 0, 0, 0], dtype=np.float32)
        re_out, im_out = fft.fft(re)
        # All real parts should be 1.0 (flat spectrum)
        np.testing.assert_array_almost_equal(re_out, [1, 1, 1, 1], decimal=3)
        np.testing.assert_array_almost_equal(im_out, [0, 0, 0, 0], decimal=3)

    def test_dc_signal(self):
        """FFT of constant [1, 1, 1, 1] = DC spike [4, 0, 0, 0]."""
        fft = SBFFT(max_size=1024)
        re = np.array([1, 1, 1, 1], dtype=np.float32)
        re_out, im_out = fft.fft(re)
        self.assertAlmostEqual(re_out[0], 4.0, places=2)
        for i in range(1, 4):
            self.assertAlmostEqual(abs(re_out[i]), 0.0, places=2)

    def test_round_trip(self):
        """FFT then inverse FFT recovers original signal."""
        fft = SBFFT(max_size=1024)
        original = np.random.randn(256).astype(np.float32)
        re_freq, im_freq = fft.fft(original)
        re_back, im_back = fft.fft(re_freq, im_freq, inverse=True)
        np.testing.assert_array_almost_equal(re_back, original, decimal=2)

    def test_sine_wave(self):
        """FFT of sine wave should peak at the correct frequency bin."""
        fft = SBFFT(max_size=1024)
        N = 256
        freq_bin = 10  # sine at bin 10
        t = np.arange(N, dtype=np.float32) / N
        signal = np.sin(2 * np.pi * freq_bin * t).astype(np.float32)
        re_out, im_out = fft.fft(signal)
        magnitudes = np.sqrt(re_out**2 + im_out**2)
        peak_bin = np.argmax(magnitudes[:N//2])
        self.assertEqual(peak_bin, freq_bin)

    def test_large_fft(self):
        """4096-point FFT round-trip."""
        fft = SBFFT(max_size=8192)
        original = np.random.randn(4096).astype(np.float32)
        re_freq, im_freq = fft.fft(original)
        re_back, _ = fft.fft(re_freq, im_freq, inverse=True)
        np.testing.assert_array_almost_equal(re_back, original, decimal=1)

    def test_parseval_theorem(self):
        """Sum of |signal|² = sum of |spectrum|² / N (Parseval's theorem)."""
        fft = SBFFT(max_size=1024)
        signal = np.random.randn(128).astype(np.float32)
        re_freq, im_freq = fft.fft(signal)
        energy_time = np.sum(signal**2)
        energy_freq = np.sum(re_freq**2 + im_freq**2) / len(signal)
        self.assertAlmostEqual(energy_time, energy_freq, places=0)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestSBFabricIntegration(unittest.TestCase):
    """Integration tests: compose primitives."""

    def test_histogram_of_scan(self):
        """Scan produces cumulative sums, histogram bins them."""
        scan = SBScan(max_size=1024)
        hist = SBHistogram(max_size=1024)

        data = np.ones(100, dtype=np.float32)
        scanned = scan.scan(data, op="sum")  # [0, 1, 2, ..., 99]
        counts = hist.histogram(scanned, num_bins=10, val_range=(0.0, 100.0))

        self.assertEqual(int(np.sum(counts)), 100)
        # Should be roughly 10 per bin (uniform distribution after scan)
        for c in counts:
            self.assertGreater(c, 5)
            self.assertLess(c, 15)

    def test_convolve_then_histogram(self):
        """Convolve smooths, histogram shows energy distribution."""
        conv = SBConvolve(max_size=4096)
        hist = SBHistogram(max_size=4096)

        data = np.random.uniform(0, 1, 1000).astype(np.float32)
        kernel = np.array([0.25, 0.5, 0.25], dtype=np.float32)
        smoothed = conv.convolve(data, kernel)

        counts_raw = hist.histogram(data, num_bins=20, val_range=(0.0, 1.0))
        counts_smooth = hist.histogram(smoothed, num_bins=20, val_range=(0.0, 1.0))

        # Both histograms should account for all elements
        self.assertEqual(int(np.sum(counts_raw)), 1000)
        self.assertEqual(int(np.sum(counts_smooth)), 1000)


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