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

Dissolves the boundaries between CPU, FPGA, and ASIC virtual layers.
All three share the SAME GPU-resident Fabric slots. Zero-copy between layers.

Before (isolated layers):
  CPU output → Python dict → download → re-upload → FPGA input
  FPGA output → Python dict → download → re-upload → ASIC input
  Cost: ~2ms per layer crossing (upload + download)

After (spectral bridge):
  CPU writes to slot 'a' → FPGA reads slot 'a' → ASIC reads slot 'a'
  All on GPU. Zero transfer between layers. One spectral pipeline.

The bridge maps:
  CPU registers    → GPU slot region [0:16]
  CPU data memory  → GPU slot region [16:272]
  FPGA block state → GPU slot region [272:528]
  ASIC hash state  → GPU slot region [528:536]

Usage:
    bridge = SpectralBridge(fabric)

    # Load and run CPU program that produces data
    bridge.cpu_load([LOAD_IMM R1, 42, STORE [0], R1, HALT])
    bridge.cpu_run(100)

    # FPGA circuit reads CPU's data memory as input
    bridge.fpga_load(blocks, connections)
    bridge.fpga_run(10)

    # ASIC hashes FPGA output
    digest = bridge.asic_hash(bridge.read_region('fpga_output', 16))

    # Or: run the whole pipeline in one call
    result = bridge.run_pipeline(cpu_program, fpga_circuit, asic_task)
"""

import numpy as np
import struct
import time
from typing import Optional


class SpectralBridge:
    """Cross-layer spectral pipeline: CPU ↔ FPGA ↔ ASIC through shared GPU slots.

    All layers share one SoundBlockFabric instance and its GPU-resident slots.
    Data flows between layers without leaving the GPU conceptual space.
    """

    # Memory map: regions within the shared address space
    # Each region is an offset into a flat array
    CPU_REGS_OFFSET = 0        # 16 registers
    CPU_REGS_SIZE = 16
    CPU_DMEM_OFFSET = 16       # 256 words data memory
    CPU_DMEM_SIZE = 256
    FPGA_STATE_OFFSET = 272    # 256 words FPGA block state
    FPGA_STATE_SIZE = 256
    ASIC_STATE_OFFSET = 528    # 8 words hash state
    ASIC_STATE_SIZE = 8
    BRIDGE_IO_OFFSET = 536     # 256 words shared I/O buffer
    BRIDGE_IO_SIZE = 256
    COLONY_WEIGHTS_OFFSET = 792   # colony concept weights (max 16 strategies × 16 dims)
    COLONY_WEIGHTS_SIZE = 256
    COLONY_MATRIX_OFFSET = 1048   # colony influence matrix (max 16×16)
    COLONY_MATRIX_SIZE = 256
    COLONY_META_OFFSET = 1304     # colony metadata (strategy count, convergence, timing)
    COLONY_META_SIZE = 32
    TOTAL_SIZE = 1336

    def __init__(self, fabric):
        self.fabric = fabric

        # Shared memory space — the bridge
        self.shared = np.zeros(self.TOTAL_SIZE, dtype=np.uint32)

        # Layer instances (lazy init)
        self._cpu = None
        self._hasher = None
        self._circuit = None

        # Stats
        self.stats = {
            'cpu_cycles': 0,
            'fpga_ticks': 0,
            'asic_hashes': 0,
            'bridge_transfers': 0,
            'total_ms': 0.0,
        }

    # ─── Shared Memory Access ────────────────────────────────────────────────

    def write_region(self, region: str, data: np.ndarray, offset: int = 0):
        """Write to a named region of shared memory."""
        base, size = self._region(region)
        n = min(len(data), size - offset)
        self.shared[base + offset:base + offset + n] = data[:n].astype(np.uint32)

    def read_region(self, region: str, count: int = None, offset: int = 0) -> np.ndarray:
        """Read from a named region of shared memory."""
        base, size = self._region(region)
        n = count if count else size
        return self.shared[base + offset:base + offset + n].copy()

    def _region(self, name: str) -> tuple:
        """Get (offset, size) for a named region."""
        regions = {
            'cpu_regs': (self.CPU_REGS_OFFSET, self.CPU_REGS_SIZE),
            'cpu_dmem': (self.CPU_DMEM_OFFSET, self.CPU_DMEM_SIZE),
            'fpga_state': (self.FPGA_STATE_OFFSET, self.FPGA_STATE_SIZE),
            'asic_state': (self.ASIC_STATE_OFFSET, self.ASIC_STATE_SIZE),
            'bridge_io': (self.BRIDGE_IO_OFFSET, self.BRIDGE_IO_SIZE),
            'colony_weights': (self.COLONY_WEIGHTS_OFFSET, self.COLONY_WEIGHTS_SIZE),
            'colony_matrix': (self.COLONY_MATRIX_OFFSET, self.COLONY_MATRIX_SIZE),
            'colony_meta': (self.COLONY_META_OFFSET, self.COLONY_META_SIZE),
        }
        if name not in regions:
            raise KeyError(f"Unknown region '{name}'. Available: {list(regions.keys())}")
        return regions[name]

    # ─── CPU Layer ───────────────────────────────────────────────────────────

    def cpu_load(self, program: list):
        """Load a program into the CPU layer."""
        from .sb_spectral_cpu import SpectralCPU
        self._cpu = SpectralCPU(self.fabric)
        self._cpu.load_program(program)

    def cpu_run(self, max_cycles: int = 1000) -> dict:
        """Run CPU and sync state to shared memory."""
        if self._cpu is None:
            raise RuntimeError("No CPU program loaded. Call cpu_load() first.")

        # Sync shared memory INTO cpu before running
        self._cpu.registers[:] = self.shared[self.CPU_REGS_OFFSET:
                                              self.CPU_REGS_OFFSET + self.CPU_REGS_SIZE]
        self._cpu.dmem[:self.CPU_DMEM_SIZE] = self.shared[self.CPU_DMEM_OFFSET:
                                                           self.CPU_DMEM_OFFSET + self.CPU_DMEM_SIZE]

        result = self._cpu.run(max_cycles)

        # Sync cpu state BACK to shared memory
        self.shared[self.CPU_REGS_OFFSET:
                    self.CPU_REGS_OFFSET + self.CPU_REGS_SIZE] = self._cpu.registers[:self.CPU_REGS_SIZE]
        self.shared[self.CPU_DMEM_OFFSET:
                    self.CPU_DMEM_OFFSET + self.CPU_DMEM_SIZE] = self._cpu.dmem[:self.CPU_DMEM_SIZE]

        self.stats['cpu_cycles'] += result['cycles']
        self.stats['bridge_transfers'] += 1
        return result

    # ─── FPGA Layer ──────────────────────────────────────────────────────────

    def fpga_load(self, blocks: list, connections: list):
        """Load an FPGA circuit."""
        self._circuit = self.fabric.create_circuit(blocks, connections)

    def fpga_run(self, n_ticks: int = 10, **kwargs) -> dict:
        """Run FPGA circuit, reading input from shared memory."""
        if self._circuit is None:
            raise RuntimeError("No FPGA circuit loaded. Call fpga_load() first.")

        # Sync shared FPGA region into circuit block states
        block_names = list(self._circuit._blocks.keys())
        for i, name in enumerate(block_names):
            if i < self.FPGA_STATE_SIZE:
                val = float(self.shared[self.FPGA_STATE_OFFSET + i])
                self._circuit._blocks[name]['state'] = np.array([val], dtype=np.float32)

        result = self._circuit.run_fused(n_ticks, **kwargs)

        # Sync circuit output BACK to shared FPGA region
        for i, name in enumerate(block_names):
            if i < self.FPGA_STATE_SIZE:
                state = self._circuit._blocks[name]['state']
                self.shared[self.FPGA_STATE_OFFSET + i] = np.uint32(int(state[0]) & 0xFFFFFFFF) \
                    if abs(state[0]) < 2**31 else 0

        self.stats['fpga_ticks'] += n_ticks
        self.stats['bridge_transfers'] += 1
        return result

    # ─── ASIC Layer (SHA-256) ────────────────────────────────────────────────

    def asic_hash(self, data: bytes) -> bytes:
        """Hash data using the spectral SHA-256 ASIC."""
        from .sb_spectral_sha256 import SpectralSHA256
        if self._hasher is None:
            self._hasher = SpectralSHA256(self.fabric)

        digest = self._hasher.hash_optimized(data)

        # Store hash result in shared ASIC region
        for i in range(8):
            self.shared[self.ASIC_STATE_OFFSET + i] = struct.unpack('>I', digest[i*4:(i+1)*4])[0]

        self.stats['asic_hashes'] += 1
        self.stats['bridge_transfers'] += 1
        return digest

    def asic_hash_batch(self, messages: list) -> list:
        """Batch hash using fused GPU shader."""
        from .sb_spectral_sha256 import SpectralSHA256
        if self._hasher is None:
            self._hasher = SpectralSHA256(self.fabric)

        digests = self._hasher.hash_fused(messages)
        self.stats['asic_hashes'] += len(messages)
        return digests

    # ─── Colony Layer (Decision Engine) ────────────────────────────────────

    def colony_load_from_bridge(self, n_dims: int):
        """Load colony influence matrix and attractors from bridge shared memory.

        Sound Block native: colony reads its configuration from the bridge
        instead of Python dicts. Pre-load the colony_matrix and colony_weights
        regions, then call this to run the colony from bridge state.

        In the full VM: a SpectralCPU program constructs the world model
        and writes it to bridge regions. Then colony_load_from_bridge() runs.
        Zero Python dict building in the hot path.
        """
        # Read influence matrix from bridge colony_matrix region
        matrix_flat = self.read_region('colony_matrix', n_dims * n_dims)
        matrix = np.zeros((n_dims, n_dims), dtype=np.float32)
        for i in range(n_dims):
            for j in range(n_dims):
                raw = int(matrix_flat[i * n_dims + j]) if i * n_dims + j < len(matrix_flat) else 0
                # Decode from fixed-point (stored as val*1000 + 0x80000000)
                matrix[i, j] = (raw - 0x80000000) / 1000.0 if raw != 0 else 0.0
        return matrix

    def colony_run(self, factors: list, relationships: list, options: list,
                   config: dict = None) -> dict:
        """Run colony decision engine, store results in shared memory.

        The colony IS the Sound Block discovery engine. It finds optimal
        strategies through FM modulation of probes through the influence matrix.
        Results flow to other layers via the bridge.

        NOTE: The Python dict construction here is the BOOTSTRAP.
        In the full Sound Block VM, a SpectralCPU program builds the world model
        from bridge I/O and writes directly to colony regions. The colony kernel
        then reads from GPU-resident data. Zero Python in the hot path.

        Args:
            factors: [{name, range}]
            relationships: [{from, to, strength}]
            options: [{name, factor: value, ...}]
            config: override colony config

        Returns:
            Colony result dict (strategies, timings, etc.)
        """
        from .sb_colony_block_kernel import SBColonyBlockKernel

        n_dims = len(factors)
        kernel = SBColonyBlockKernel(num_dims=n_dims)

        # Bootstrap: Python dict construction (will become SpectralCPU program)
        input_dict = {
            'factors': factors,
            'relationships': relationships,
            'options': options,
        }

        result = kernel.run(input_dict, config)

        # Store concept weights in shared colony region
        strategies = result.get('strategies', [])
        for si, strat in enumerate(strategies):
            if si >= 16:
                break
            weights = strat.get('concept_weights', {})
            for fi, f in enumerate(factors):
                if fi >= 16:
                    break
                w = weights.get(f['name'], 0.0)
                idx = self.COLONY_WEIGHTS_OFFSET + si * 16 + fi
                if idx < len(self.shared):
                    # Store as fixed-point: weight × 1000
                    self.shared[idx] = np.uint32(int(w * 1000) & 0xFFFFFFFF)

        # Store influence matrix in shared colony region
        inf_matrix = result.get('influence_matrix', [])
        for i, row in enumerate(inf_matrix):
            if i >= 16:
                break
            for j, val in enumerate(row):
                if j >= 16:
                    break
                idx = self.COLONY_MATRIX_OFFSET + i * 16 + j
                if idx < len(self.shared):
                    # Fixed-point: value × 1000
                    self.shared[idx] = np.uint32(int(val * 1000 + 0x80000000) & 0xFFFFFFFF)

        # Store metadata
        meta_base = self.COLONY_META_OFFSET
        self.shared[meta_base] = len(strategies)
        self.shared[meta_base + 1] = n_dims
        self.shared[meta_base + 2] = result.get('pass_count', 0)
        self.shared[meta_base + 3] = 1 if result.get('converged', False) else 0

        self.stats['bridge_transfers'] += 1
        return result

    def colony_to_fpga(self, strategy_index: int = 0):
        """Copy colony strategy weights to FPGA block initial state.

        The bridge: colony discovers optimal configuration → FPGA circuit
        executes it. Colony concept weights become FPGA block amplitudes.
        """
        base = self.COLONY_WEIGHTS_OFFSET + strategy_index * 16
        n_dims = int(self.shared[self.COLONY_META_OFFSET + 1])
        for i in range(min(n_dims, self.FPGA_STATE_SIZE)):
            w = self.shared[base + i]
            self.shared[self.FPGA_STATE_OFFSET + i] = w
        self.stats['bridge_transfers'] += 1

    def colony_to_cpu(self, strategy_index: int = 0, dmem_offset: int = 0):
        """Copy colony strategy weights to CPU data memory.

        The bridge: colony output → CPU program input.
        CPU can read weights and make decisions based on colony results.
        """
        base = self.COLONY_WEIGHTS_OFFSET + strategy_index * 16
        n_dims = int(self.shared[self.COLONY_META_OFFSET + 1])
        for i in range(min(n_dims, 16)):
            self.shared[self.CPU_DMEM_OFFSET + dmem_offset + i] = self.shared[base + i]
        self.stats['bridge_transfers'] += 1

    def colony_to_asic(self, strategy_index: int = 0) -> bytes:
        """Hash colony strategy output for integrity verification.

        The bridge: colony output → ASIC SHA-256 hash.
        Proves the strategy hasn't been tampered with.
        """
        base = self.COLONY_WEIGHTS_OFFSET + strategy_index * 16
        n_dims = int(self.shared[self.COLONY_META_OFFSET + 1])
        data = b''.join(
            struct.pack('>I', int(self.shared[base + i]))
            for i in range(min(n_dims, 16))
        )
        return self.asic_hash(data)

    # ─── Cross-Layer Operations ──────────────────────────────────────────────

    def cpu_to_fpga(self, cpu_dmem_offset: int = 0, count: int = 16):
        """Copy CPU data memory region to FPGA input state.

        The bridge: CPU output → FPGA input through shared memory.
        No GPU transfer — just array slice within shared space.
        """
        src = self.CPU_DMEM_OFFSET + cpu_dmem_offset
        dst = self.FPGA_STATE_OFFSET
        n = min(count, self.FPGA_STATE_SIZE)
        self.shared[dst:dst + n] = self.shared[src:src + n]
        self.stats['bridge_transfers'] += 1

    def fpga_to_asic(self, count: int = 8):
        """Copy FPGA output state to ASIC input region.

        The bridge: FPGA output → ASIC input through shared memory.
        """
        src = self.FPGA_STATE_OFFSET
        dst = self.BRIDGE_IO_OFFSET
        n = min(count, self.BRIDGE_IO_SIZE)
        self.shared[dst:dst + n] = self.shared[src:src + n]
        self.stats['bridge_transfers'] += 1

    def cpu_to_asic(self, cpu_dmem_offset: int = 0, byte_count: int = 32) -> bytes:
        """Read CPU data memory as bytes and hash through ASIC.

        The full bridge: CPU produces data → ASIC hashes it.
        """
        words = self.shared[self.CPU_DMEM_OFFSET + cpu_dmem_offset:
                            self.CPU_DMEM_OFFSET + cpu_dmem_offset + (byte_count + 3) // 4]
        data = b''.join(struct.pack('>I', int(w)) for w in words)[:byte_count]
        return self.asic_hash(data)

    def asic_to_cpu(self, cpu_dmem_offset: int = 0):
        """Copy ASIC hash result into CPU data memory.

        The bridge: ASIC output → CPU input for verification.
        """
        for i in range(8):
            self.shared[self.CPU_DMEM_OFFSET + cpu_dmem_offset + i] = \
                self.shared[self.ASIC_STATE_OFFSET + i]
        self.stats['bridge_transfers'] += 1

    # ─── GPU Fabric Layer (Direct Compute) ─────────────────────────────────

    def gpu_upload_region(self, region: str, slot: str = 'a'):
        """Upload a shared memory region to a GPU-resident slot.

        The GPU bridge: shared memory → GPU slot for Fabric compute.
        After this, all 26 Fabric primitives operate on the data
        (matmul, FFT, map, convolve, etc.) without CPU involvement.
        """
        base, size = self._region(region)
        data = self.shared[base:base + size].astype(np.float32)
        self.fabric.upload(slot, data)
        self.fabric._slot_sizes[slot] = size
        self.stats['bridge_transfers'] += 1

    def gpu_download_region(self, slot: str, region: str):
        """Download a GPU-resident slot back to shared memory region.

        The GPU bridge: GPU slot → shared memory for cross-layer access.
        """
        base, size = self._region(region)
        data = self.fabric.download(slot)[:size]
        # Convert float32 back to uint32 for shared memory
        u32_data = np.frombuffer(data.astype(np.float32).tobytes(),
                                 dtype=np.uint32)[:size]
        self.shared[base:base + min(len(u32_data), size)] = u32_data[:size]
        self.stats['bridge_transfers'] += 1

    def gpu_compute(self, src_region: str, dst_region: str,
                    operations: list, src_slot: str = 'a', dst_slot: str = 'out'):
        """Run a chain of GPU Fabric operations on a shared memory region.

        The GPU bridge: upload → chain ops on GPU → download.
        Operations: list of (op_name, args) tuples.

        Example:
            bridge.gpu_compute('cpu_dmem', 'bridge_io', [
                ('map', 'tanh'),
                ('map', 'scale', 2.0),
            ])
        """
        self.gpu_upload_region(src_region, src_slot)

        current = src_slot
        other = dst_slot
        for op in operations:
            name = op[0]
            if name == 'map':
                op_name = op[1] if len(op) > 1 else 'abs'
                scalar = op[2] if len(op) > 2 else 0.0
                self.fabric.map_gpu(current, other, op_name, scalar)
            elif name == 'add':
                # Need data in both slots
                self.fabric.add_gpu(current, op[1] if len(op) > 1 else 'b', other)
            elif name == 'xor':
                self.fabric.xor_gpu(current, op[1] if len(op) > 1 else 'b', other)
            elif name == 'matmul':
                M, K, N = op[1], op[2], op[3]
                self.fabric.matmul_gpu(current, 'mat_a', other, M, K, N)
            else:
                continue
            # Swap slots for chaining
            current, other = other, current

        self.gpu_download_region(current, dst_region)

    def gpu_spectral_transform(self, region: str, inverse: bool = False):
        """Apply FFT to a shared memory region — the spectral transform.

        This IS the bridge between discrete and spectral domains.
        Transforms time-domain data into frequency-domain for spectral processing.
        """
        base, size = self._region(region)
        data = self.shared[base:base + size].astype(np.float32)

        # Pad to power of 2 for FFT
        n = 1
        while n < len(data):
            n *= 2
        padded = np.zeros(n, dtype=np.float32)
        padded[:len(data)] = data

        re, im = self.fabric.fft(padded, inverse=inverse)

        # Write back (magnitudes for forward, values for inverse)
        if not inverse:
            magnitudes = np.sqrt(re**2 + im**2)[:size]
            result = np.frombuffer(magnitudes.astype(np.float32).tobytes(),
                                   dtype=np.uint32)
        else:
            result = np.frombuffer(re[:size].astype(np.float32).tobytes(),
                                   dtype=np.uint32)
        self.shared[base:base + min(len(result), size)] = result[:size]
        self.stats['bridge_transfers'] += 1

    # ─── Full Pipeline ───────────────────────────────────────────────────────

    def run_pipeline(self, cpu_program: list, cpu_cycles: int = 1000,
                     fpga_blocks: list = None, fpga_connections: list = None,
                     fpga_ticks: int = 10,
                     hash_cpu_output: bool = False,
                     hash_offset: int = 0, hash_bytes: int = 32) -> dict:
        """Run the complete cross-layer pipeline.

        CPU → (optional) FPGA → (optional) ASIC hash of output.
        All layers share the same memory space. Zero-copy between layers.

        Returns dict with all layer results + bridge stats.
        """
        t0 = time.perf_counter()

        # Stage 1: CPU
        self.cpu_load(cpu_program)
        cpu_result = self.cpu_run(cpu_cycles)

        # Stage 2: FPGA (if circuit provided)
        fpga_result = None
        if fpga_blocks:
            self.fpga_load(fpga_blocks, fpga_connections or [])
            self.cpu_to_fpga()
            fpga_result = self.fpga_run(fpga_ticks)

        # Stage 3: ASIC hash (if requested)
        digest = None
        if hash_cpu_output:
            digest = self.cpu_to_asic(hash_offset, hash_bytes)

        total_ms = (time.perf_counter() - t0) * 1000
        self.stats['total_ms'] += total_ms

        return {
            'cpu': cpu_result,
            'fpga': fpga_result,
            'digest': digest.hex() if digest else None,
            'shared_memory': self.shared.copy(),
            'stats': self.stats.copy(),
            'total_ms': total_ms,
        }
