"""
Spectral Bridge — Cross-Layer Unified Pipeline Tests.

Tests CPU ↔ FPGA ↔ ASIC ↔ GPU bridges through shared memory.
Like Apple's unified memory architecture, but for compute types:
all processing shares one spectral fabric.

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

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

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

from afs.gpu.sb_fabric import SoundBlockFabric, gpu_available
from afs.gpu.sb_spectral_bridge import SpectralBridge
from afs.gpu.sb_spectral_cpu import SpectralCPU


E = SpectralCPU.encode


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestBridgeSharedMemory(unittest.TestCase):
    """Shared memory: zero-copy between layers."""

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

    def test_write_read_region(self):
        """Write to CPU dmem, read back."""
        data = np.array([10, 20, 30, 40], dtype=np.uint32)
        self.bridge.write_region('cpu_dmem', data)
        result = self.bridge.read_region('cpu_dmem', 4)
        np.testing.assert_array_equal(result, data)

    def test_cross_region_isolation(self):
        """Writing to one region doesn't affect another."""
        self.bridge.write_region('cpu_regs', np.array([99], dtype=np.uint32))
        self.bridge.write_region('fpga_state', np.array([42], dtype=np.uint32))
        self.assertEqual(self.bridge.read_region('cpu_regs', 1)[0], 99)
        self.assertEqual(self.bridge.read_region('fpga_state', 1)[0], 42)

    def test_cpu_to_fpga_bridge(self):
        """CPU data memory → FPGA state through shared memory."""
        self.bridge.write_region('cpu_dmem', np.array([100, 200, 300], dtype=np.uint32))
        self.bridge.cpu_to_fpga(count=3)
        fpga_data = self.bridge.read_region('fpga_state', 3)
        np.testing.assert_array_equal(fpga_data, [100, 200, 300])
        print(f"\n  CPU→FPGA bridge: [100,200,300] ✓")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestBridgeCPU(unittest.TestCase):
    """CPU layer through the bridge."""

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

    def test_cpu_computes_through_bridge(self):
        """CPU runs program, results visible in shared memory."""
        self.bridge.cpu_load([
            E(0x8, 1, imm=42),     # R1 = 42
            E(0x8, 2, imm=10),     # R2 = 10
            E(0xA, 0, 2, 1, imm=0),# STORE [R2+0], R1 → dmem[10] = 42
            E(0xF),
        ])
        self.bridge.cpu_run(100)

        # Check shared memory directly
        regs = self.bridge.read_region('cpu_regs', 16)
        dmem = self.bridge.read_region('cpu_dmem', 256)
        self.assertEqual(regs[1], 42)
        self.assertEqual(dmem[10], 42)

    def test_cpu_reads_preset_shared_memory(self):
        """Pre-load shared memory, CPU reads it."""
        # Pre-load data into CPU dmem region
        self.bridge.write_region('cpu_dmem', np.array([7, 8, 9], dtype=np.uint32))

        self.bridge.cpu_load([
            E(0x8, 1, imm=0),      # R1 = 0 (address)
            E(0x9, 2, 1, 0, imm=0),# LOAD R2, [R1+0] → R2 = 7
            E(0x9, 3, 1, 0, imm=1),# LOAD R3, [R1+1] → R3 = 8
            E(0x1, 4, 2, 3),       # ADD R4, R2, R3  → R4 = 15
            E(0xF),
        ])
        self.bridge.cpu_run(100)

        regs = self.bridge.read_region('cpu_regs', 16)
        self.assertEqual(regs[4], 15)
        print(f"\n  CPU reads shared mem: 7+8={regs[4]} ✓")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestBridgeASIC(unittest.TestCase):
    """ASIC (SHA-256) layer through the bridge."""

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

    def test_asic_hash_matches_hashlib(self):
        """ASIC hash through bridge matches hashlib."""
        msg = b"bridge test"
        expected = hashlib.sha256(msg).digest()
        result = self.bridge.asic_hash(msg)
        self.assertEqual(result, expected)

        # Hash should be in shared ASIC region
        asic_state = self.bridge.read_region('asic_state', 8)
        self.assertGreater(asic_state[0], 0)  # non-zero hash

    def test_cpu_to_asic_bridge(self):
        """CPU produces data → ASIC hashes it. Full bridge."""
        # CPU computes and stores in dmem
        self.bridge.cpu_load([
            E(0x8, 1, imm=0x6865),  # 'he' in big-endian-ish
            E(0x8, 2, imm=0),
            E(0xA, 0, 2, 1, imm=0), # STORE [0], R1
            E(0xF),
        ])
        self.bridge.cpu_run(100)

        # ASIC hashes CPU's data memory
        digest = self.bridge.cpu_to_asic(cpu_dmem_offset=0, byte_count=4)
        self.assertEqual(len(digest), 32)  # valid SHA-256 digest

        # Result is in shared ASIC state
        asic = self.bridge.read_region('asic_state', 8)
        self.assertGreater(asic[0], 0)
        print(f"\n  CPU→ASIC bridge: hash={digest.hex()[:16]}... ✓")

    def test_asic_to_cpu_bridge(self):
        """ASIC hashes → result copied to CPU data memory."""
        self.bridge.asic_hash(b"round trip test")
        self.bridge.asic_to_cpu(cpu_dmem_offset=0)

        # CPU can now read the hash from its dmem
        dmem = self.bridge.read_region('cpu_dmem', 8)
        self.assertGreater(dmem[0], 0)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestBridgeGPU(unittest.TestCase):
    """GPU Fabric layer through the bridge."""

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

    def test_gpu_compute_on_shared_region(self):
        """GPU operates on shared memory region via Fabric primitives."""
        # Write data to bridge I/O
        data = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.uint32)
        self.bridge.write_region('bridge_io', data)

        # Upload to GPU, apply tanh, download back
        self.bridge.gpu_upload_region('bridge_io', 'a')
        self.fabric.map_gpu('a', 'out', 'scale', scalar=2.0)
        self.bridge.gpu_download_region('out', 'bridge_io')

        result = self.bridge.read_region('bridge_io', 8)
        # Values were scaled by 2 (as float32, then cast back)
        self.assertGreater(result[0], 0)
        print(f"\n  GPU compute on bridge region: {result[:4]} ✓")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestBridgeFullPipeline(unittest.TestCase):
    """Full cross-layer pipeline: CPU → FPGA → ASIC → GPU."""

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

    def test_cpu_produces_asic_hashes(self):
        """Pipeline: CPU computes data → ASIC hashes it."""
        result = self.bridge.run_pipeline(
            cpu_program=[
                E(0x8, 1, imm=42),
                E(0x8, 2, imm=0),
                E(0xA, 0, 2, 1, imm=0),
                E(0xF),
            ],
            cpu_cycles=100,
            hash_cpu_output=True,
            hash_offset=0,
            hash_bytes=4,
        )

        self.assertIsNotNone(result['digest'])
        self.assertEqual(len(result['digest']), 64)  # hex digest
        self.assertTrue(result['cpu']['halted'])
        print(f"\n  Full pipeline: CPU→ASIC hash={result['digest'][:16]}...")
        print(f"    CPU cycles: {result['cpu']['cycles']}")
        print(f"    Bridge transfers: {result['stats']['bridge_transfers']}")
        print(f"    Total: {result['total_ms']:.2f}ms")

    def test_pipeline_with_fpga(self):
        """Pipeline: CPU → FPGA → readback."""
        def counter_fn(inp, state):
            return state + 1.0, state + 1.0

        result = self.bridge.run_pipeline(
            cpu_program=[
                E(0x8, 1, imm=100),    # Store initial values for FPGA
                E(0x8, 2, imm=0),
                E(0xA, 0, 2, 1, imm=0),
                E(0xF),
            ],
            cpu_cycles=100,
            fpga_blocks=[
                {'name': 'ctr', 'fn': counter_fn, 'state_size': 1, 'gpu_type': 'counter'},
            ],
            fpga_connections=[],
            fpga_ticks=50,
        )

        self.assertTrue(result['cpu']['halted'])
        self.assertIsNotNone(result['fpga'])
        # FPGA counter ran 50 ticks
        fpga_state = self.bridge.read_region('fpga_state', 1)
        print(f"\n  CPU→FPGA pipeline: FPGA state after 50 ticks = {fpga_state[0]}")
        print(f"    Total: {result['total_ms']:.2f}ms")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestBridgeColony(unittest.TestCase):
    """Colony engine through the SpectralBridge."""

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

    def test_colony_run_produces_strategies(self):
        """Colony runs through bridge, results in shared memory."""
        factors = [
            {"name": "speed", "range": [0, 100]},
            {"name": "cost", "range": [0, 100]},
            {"name": "quality", "range": [0, 100]},
        ]
        relationships = [
            {"from": "speed", "to": "cost", "strength": 0.6},
            {"from": "cost", "to": "quality", "strength": -0.5},
        ]
        options = [
            {"name": "A", "speed": 90, "cost": 30, "quality": 70},
            {"name": "B", "speed": 40, "cost": 80, "quality": 90},
        ]

        result = self.bridge.colony_run(factors, relationships, options)
        self.assertIn('strategies', result)
        self.assertGreater(len(result['strategies']), 0)

        # Check shared memory has colony data
        meta = self.bridge.read_region('colony_meta', 4)
        self.assertGreater(meta[0], 0, "Strategy count should be > 0")
        self.assertEqual(meta[1], 3, "Dimensions should be 3")
        print(f"\n  Colony via bridge: {len(result['strategies'])} strategies, "
              f"{result['pass_count']} passes, {result['total_ms']:.1f}ms")

    def test_colony_to_cpu(self):
        """Colony results flow to CPU data memory."""
        factors = [
            {"name": "a", "range": [0, 100]},
            {"name": "b", "range": [0, 100]},
        ]
        result = self.bridge.colony_run(factors, [], [
            {"name": "X", "a": 80, "b": 20},
            {"name": "Y", "a": 20, "b": 80},
        ])

        self.bridge.colony_to_cpu(strategy_index=0, dmem_offset=0)
        dmem = self.bridge.read_region('cpu_dmem', 4)
        # Weights should be non-zero (stored as fixed-point × 1000)
        has_nonzero = any(dmem[i] != 0 for i in range(2))
        self.assertTrue(has_nonzero, "Colony weights should flow to CPU dmem")
        print(f"\n  Colony→CPU: weights in dmem = {dmem[:2]}")

    def test_colony_to_asic_hash(self):
        """Colony output hashed for integrity via ASIC bridge."""
        factors = [
            {"name": "x", "range": [0, 100]},
            {"name": "y", "range": [0, 100]},
        ]
        self.bridge.colony_run(factors, [], [
            {"name": "A", "x": 50, "y": 50},
        ])

        digest = self.bridge.colony_to_asic(strategy_index=0)
        self.assertEqual(len(digest), 32)
        print(f"\n  Colony→ASIC hash: {digest.hex()[:16]}...")

    def test_full_colony_pipeline(self):
        """Colony → CPU → ASIC: complete cross-layer flow."""
        factors = [
            {"name": "price", "range": [0, 100]},
            {"name": "quality", "range": [0, 100]},
        ]
        relationships = [{"from": "price", "to": "quality", "strength": -0.5}]
        options = [
            {"name": "Budget", "price": 20, "quality": 40},
            {"name": "Premium", "price": 80, "quality": 90},
        ]

        # Colony discovers strategies
        result = self.bridge.colony_run(factors, relationships, options)

        # Flow to CPU
        self.bridge.colony_to_cpu(strategy_index=0)

        # CPU loads weights, does arithmetic
        self.bridge.cpu_load([
            E(0x8, 1, imm=0),       # R1 = 0 (address)
            E(0x9, 2, 1, 0, imm=0), # LOAD R2, [0] (price weight)
            E(0x9, 3, 1, 0, imm=1), # LOAD R3, [1] (quality weight)
            E(0x1, 4, 2, 3),        # R4 = price + quality weight
            E(0xF),
        ])
        self.bridge.cpu_run(100)

        # Hash the CPU result
        digest = self.bridge.cpu_to_asic(cpu_dmem_offset=0, byte_count=8)

        regs = self.bridge.read_region('cpu_regs', 16)
        print(f"\n  Full pipeline: Colony→CPU→ASIC")
        print(f"    Colony: {len(result['strategies'])} strategies")
        print(f"    CPU R4 (sum of weights): {regs[4]}")
        print(f"    ASIC hash: {digest.hex()[:16]}...")
        print(f"    Bridge transfers: {self.bridge.stats['bridge_transfers']}")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestBridgeBenchmark(unittest.TestCase):
    """Performance: how fast is cross-layer communication?"""

    def test_bridge_overhead(self):
        """Measure the cost of cross-layer bridge transfers."""
        fabric = SoundBlockFabric(max_size=10240)
        bridge = SpectralBridge(fabric)

        # Measure CPU→FPGA bridge cost
        bridge.write_region('cpu_dmem', np.arange(256, dtype=np.uint32))
        t0 = time.perf_counter()
        for _ in range(10000):
            bridge.cpu_to_fpga(count=16)
        bridge_ms = (time.perf_counter() - t0) * 1000

        # Measure GPU upload/download cost
        data = np.random.randn(256).astype(np.float32)
        t0 = time.perf_counter()
        for _ in range(100):
            fabric.upload('a', data)
            fabric.download('a')
        gpu_ms = (time.perf_counter() - t0) * 1000

        print(f"\n  Bridge Overhead Benchmark:")
        print(f"    CPU→FPGA bridge (16 words): {bridge_ms/10000*1000:.1f}us per transfer")
        print(f"    GPU upload+download (256 floats): {gpu_ms/100*1000:.1f}us per round-trip")
        print(f"    Bridge is shared memory (array slice) — near-zero cost")
        print(f"    GPU transfer is actual device I/O — ~{gpu_ms/100:.2f}ms")

    def test_full_pipeline_speed(self):
        """End-to-end pipeline throughput."""
        fabric = SoundBlockFabric(max_size=10240)
        bridge = SpectralBridge(fabric)

        program = [
            E(0x8, 1, imm=42),
            E(0x8, 2, imm=0),
            E(0xA, 0, 2, 1, imm=0),
            E(0xF),
        ]

        # Warmup
        bridge.run_pipeline(program, cpu_cycles=100,
                           hash_cpu_output=True, hash_bytes=4)

        times = []
        for _ in range(20):
            b = SpectralBridge(fabric)
            t0 = time.perf_counter()
            b.run_pipeline(program, cpu_cycles=100,
                          hash_cpu_output=True, hash_bytes=4)
            times.append((time.perf_counter() - t0) * 1000)

        avg = np.mean(times)
        print(f"\n  Full Pipeline (CPU→ASIC) Benchmark:")
        print(f"    Mean: {avg:.2f}ms ± {np.std(times):.2f}ms")
        print(f"    Throughput: {1000/avg:.0f} pipelines/sec")


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