"""
Spectral CPU — Tests for the Sound Block virtual CPU processor.

Every CPU pipeline stage is a Sound Block operation:
  Fetch = memory read, Decode = LUT, Execute = ALU primitives,
  Writeback = register write, Branch = condition LUT + PC update.

Run:
    cd afs-engine
    python3 tests/test_sb_spectral_cpu.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
from afs.gpu.sb_spectral_cpu import SpectralCPU


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestCPUArithmetic(unittest.TestCase):
    """Test all ALU operations."""

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

    def test_add(self):
        """ADD R3, R1, R2."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=10),     # LOAD_IMM R1, 10
            E(0x8, 2, imm=20),     # LOAD_IMM R2, 20
            E(0x1, 3, 1, 2),       # ADD R3, R1, R2
            E(0xF),                 # HALT
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 30)

    def test_sub(self):
        """SUB R3, R1, R2."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=50),
            E(0x8, 2, imm=20),
            E(0x2, 3, 1, 2),       # SUB R3, R1, R2
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 30)

    def test_and(self):
        """AND R3, R1, R2."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=0xFF00),
            E(0x8, 2, imm=0x0FF0),
            E(0x3, 3, 1, 2),
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 0x0F00)

    def test_or(self):
        """OR R3, R1, R2."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=0xF000),
            E(0x8, 2, imm=0x00F0),
            E(0x4, 3, 1, 2),
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 0xF0F0)

    def test_xor(self):
        """XOR R3, R1, R2."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=0xFF00),
            E(0x8, 2, imm=0xF0F0),
            E(0x5, 3, 1, 2),
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 0x0FF0)

    def test_shr(self):
        """SHR R3, R1, R2."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=0xFF00),
            E(0x8, 2, imm=4),
            E(0x6, 3, 1, 2),
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 0x0FF0)

    def test_mul(self):
        """MUL R3, R1, R2."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=7),
            E(0x8, 2, imm=6),
            E(0xE, 3, 1, 2),       # MUL
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 42)

    def test_load_imm(self):
        """LOAD_IMM R5, 12345."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 5, imm=12345),
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[5], 12345)

    def test_r0_always_zero(self):
        """R0 cannot be written (RISC convention)."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 0, imm=999),    # Try to write R0
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[0], 0)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestCPUMemory(unittest.TestCase):
    """Test load/store operations."""

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

    def test_store_and_load(self):
        """Store a value then load it back."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=42),     # R1 = 42
            E(0x8, 2, imm=10),     # R2 = 10 (address)
            E(0xA, 0, 2, 1, imm=0),# STORE [R2+0], R1
            E(0x9, 3, 2, 0, imm=0),# LOAD R3, [R2+0]
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 42)
        self.assertEqual(self.cpu.dmem[10], 42)

    def test_store_with_offset(self):
        """Store with immediate offset."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=99),
            E(0x8, 2, imm=0),      # base = 0
            E(0xA, 0, 2, 1, imm=5),# STORE [R2+5], R1
            E(0x9, 3, 2, 0, imm=5),# LOAD R3, [R2+5]
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 99)
        self.assertEqual(self.cpu.dmem[5], 99)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestCPUBranch(unittest.TestCase):
    """Test branch instructions."""

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

    def test_beq_taken(self):
        """BEQ when equal — should skip instruction."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=5),      # R1 = 5
            E(0x8, 2, imm=5),      # R2 = 5
            E(0xB, 0, 1, 2, imm=2),# BEQ R1, R2, +2 (skip next)
            E(0x8, 3, imm=999),    # R3 = 999 (should be skipped)
            E(0x8, 4, imm=1),      # R4 = 1 (should execute)
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 0)   # skipped
        self.assertEqual(self.cpu.registers[4], 1)    # executed

    def test_beq_not_taken(self):
        """BEQ when not equal — should NOT skip."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=5),
            E(0x8, 2, imm=10),
            E(0xB, 0, 1, 2, imm=2),# BEQ R1, R2, +2 (NOT taken)
            E(0x8, 3, imm=999),    # R3 = 999 (should execute)
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 999)

    def test_bne(self):
        """BNE — branch when not equal."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=5),
            E(0x8, 2, imm=10),
            E(0xC, 0, 1, 2, imm=2),# BNE R1, R2, +2 (taken)
            E(0x8, 3, imm=999),    # skipped
            E(0x8, 4, imm=1),
            E(0xF),
        ])
        self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 0)
        self.assertEqual(self.cpu.registers[4], 1)

    def test_loop(self):
        """Count from 0 to 10 using a loop."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=0),      # R1 = 0 (counter)
            E(0x8, 2, imm=1),      # R2 = 1 (increment)
            E(0x8, 3, imm=10),     # R3 = 10 (limit)
            # loop:
            E(0x1, 1, 1, 2),       # ADD R1, R1, R2  (counter++)
            E(0xD, 0, 1, 3, imm=0xFFFE),# BLT R1, R3, -2 (back to ADD)
            E(0xF),                 # HALT
        ])
        result = self.cpu.run()
        self.assertEqual(self.cpu.registers[1], 10)
        print(f"\n  Loop 0→10: {result['cycles']} cycles, R1={self.cpu.registers[1]}")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestCPUPrograms(unittest.TestCase):
    """Test real programs on the Spectral CPU."""

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

    def test_fibonacci(self):
        """Compute Fibonacci numbers: fib(10) = 55."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=0),      # R1 = fib(n-2) = 0
            E(0x8, 2, imm=1),      # R2 = fib(n-1) = 1
            E(0x8, 3, imm=9),      # R3 = count = 9 (fib(0),fib(1) already set)
            E(0x8, 4, imm=1),      # R4 = 1 (decrement)
            E(0x8, 5, imm=0),      # R5 = 0 (for comparison)
            # loop:
            E(0x1, 6, 1, 2),       # R6 = R1 + R2 (next fib)
            E(0x1, 1, 2, 5),       # R1 = R2 (shift: add R2 + 0)
            # Actually: R1 = R2. Use ADD R1, R2, R5 (R5=0)
            E(0x1, 2, 6, 5),       # R2 = R6 (ADD R2, R6, R5)
            E(0x2, 3, 3, 4),       # R3 = R3 - 1
            E(0xC, 0, 3, 5, imm=0xFFFB),# BNE R3, R5, -5 (back to loop)
            E(0xF),                 # HALT
        ])
        result = self.cpu.run()
        self.assertEqual(self.cpu.registers[2], 55,
            f"fib(10) should be 55, got {self.cpu.registers[2]}")
        print(f"\n  Fibonacci(10) = {self.cpu.registers[2]}, "
              f"{result['cycles']} cycles")

    def test_sum_array(self):
        """Sum an array stored in data memory."""
        E = SpectralCPU.encode
        # Store array [10, 20, 30, 40, 50] in dmem[0..4]
        self.cpu.dmem[0] = 10
        self.cpu.dmem[1] = 20
        self.cpu.dmem[2] = 30
        self.cpu.dmem[3] = 40
        self.cpu.dmem[4] = 50

        self.cpu.load_program([
            E(0x8, 1, imm=0),      # R1 = 0 (sum)
            E(0x8, 2, imm=0),      # R2 = 0 (index)
            E(0x8, 3, imm=5),      # R3 = 5 (count)
            E(0x8, 4, imm=1),      # R4 = 1 (increment)
            # loop:
            E(0x9, 5, 2, 0, imm=0),# LOAD R5, [R2+0]
            E(0x1, 1, 1, 5),       # R1 += R5
            E(0x1, 2, 2, 4),       # R2 += 1
            E(0xD, 0, 2, 3, imm=0xFFFC),# BLT R2, R3, -4
            E(0xF),
        ])
        result = self.cpu.run()
        self.assertEqual(self.cpu.registers[1], 150)
        print(f"\n  Sum [10,20,30,40,50] = {self.cpu.registers[1]}, "
              f"{result['cycles']} cycles")

    def test_multiply_by_addition(self):
        """Multiply 7 × 8 using repeated addition."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=7),      # R1 = 7 (multiplicand)
            E(0x8, 2, imm=8),      # R2 = 8 (multiplier/counter)
            E(0x8, 3, imm=0),      # R3 = 0 (result)
            E(0x8, 4, imm=1),      # R4 = 1
            E(0x8, 5, imm=0),      # R5 = 0
            # loop:
            E(0x1, 3, 3, 1),       # R3 += R1
            E(0x2, 2, 2, 4),       # R2 -= 1
            E(0xC, 0, 2, 5, imm=0xFFFD),# BNE R2, R5, -3
            E(0xF),
        ])
        result = self.cpu.run()
        self.assertEqual(self.cpu.registers[3], 56)

    def test_disassemble(self):
        """Disassembly produces readable output."""
        E = SpectralCPU.encode
        self.cpu.load_program([
            E(0x8, 1, imm=42),
            E(0x1, 3, 1, 2),
            E(0xA, 0, 1, 2, imm=5),
            E(0xB, 0, 1, 2, imm=3),
            E(0xF),
        ])
        lines = self.cpu.disassemble()
        self.assertGreater(len(lines), 0)
        self.assertIn('LOAD_IMM', lines[0])
        self.assertIn('ADD', lines[1])
        print(f"\n  Disassembly:")
        for line in lines:
            print(f"  {line}")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestCPUBenchmark(unittest.TestCase):
    """Performance benchmarks for the Spectral CPU."""

    def test_cycle_speed(self):
        """Measure cycles per second."""
        fabric = SoundBlockFabric(max_size=10240)
        cpu = SpectralCPU(fabric)
        E = SpectralCPU.encode

        # Simple loop: count to 10000
        cpu.load_program([
            E(0x8, 1, imm=0),
            E(0x8, 2, imm=1),
            E(0x8, 3, imm=10000),
            E(0x1, 1, 1, 2),
            E(0xD, 0, 1, 3, imm=0xFFFE),
            E(0xF),
        ])
        result = cpu.run(max_cycles=50000)
        rate = result['instructions'] / (result['total_ms'] / 1000)

        print(f"\n  Spectral CPU Benchmark (count to 10000):")
        print(f"    Cycles: {result['cycles']}")
        print(f"    Time: {result['total_ms']:.1f}ms")
        print(f"    Rate: {rate:,.0f} instructions/sec")
        print(f"    Per-cycle: {result['per_cycle_ns']:.0f}ns")

    def test_parallel_cpus(self):
        """N virtual CPUs running in parallel."""
        fabric = SoundBlockFabric(max_size=10240)
        cpu = SpectralCPU(fabric)
        E = SpectralCPU.encode

        # Each CPU computes a different multiplication
        programs = []
        for i in range(100):
            programs.append([
                E(0x8, 1, imm=i + 1),
                E(0x8, 2, imm=i + 1),
                E(0xE, 3, 1, 2),       # MUL R3 = (i+1)²
                E(0xF),
            ])

        t0 = time.perf_counter()
        results = cpu.run_parallel(programs, max_cycles=100)
        total_ms = (time.perf_counter() - t0) * 1000

        # Verify: each CPU computed (i+1)²
        for i, r in enumerate(results):
            expected = (i + 1) ** 2
            self.assertEqual(r['registers'][3], expected,
                f"CPU {i}: expected {expected}, got {r['registers'][3]}")

        print(f"\n  Parallel CPU Benchmark (100 virtual CPUs):")
        print(f"    Total: {total_ms:.1f}ms")
        print(f"    Per-CPU: {total_ms/100:.2f}ms")
        print(f"    All 100 computed (i+1)² correctly ✓")


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