import numpy as np
import json
import time
import math
import hashlib
from datetime import datetime, timezone

Q16_SCALE = 65536
EPSILON = 1e-14


def sidon_address(k: int) -> int:
    """Sidon address for strand k: 2^k.  For 8 strands: {1,2,4,8,16,32,64,128}."""
    return 1 << k


class HouseholderReflector:
    """A Householder reflector = one braid crossing in the VCN pipeline."""
    def __init__(self, k: int, v: list[float], tau: float):
        self.k = k          # column being zeroed = VCN strand index
        self.v = v          # reflector vector
        self.tau = tau      # scaling factor
        self.sidon = sidon_address(k)


def _householder_vector(x: list[float]) -> tuple[list[float], float]:
    """Householder vector v and τ such that (I - τ·v·vᵀ)·x = μ·e₀."""
    n = len(x)
    alpha = x[0]
    nx = math.sqrt(sum(xi * xi for xi in x))
    mu = -math.copysign(nx, alpha)
    if abs(mu - alpha) < EPSILON:
        return [0.0] * n, 0.0
    v = [0.0] * n
    v[0] = 1.0
    for i in range(1, n):
        v[i] = x[i] / (alpha - mu)
    tau = (mu - alpha) / mu
    return v, tau


class VCNDSPPipeline:
    def __init__(self, num_strands=8, grid_size=8):
        """
        Initializes the VCN DSP SIMD pipeline.
        Instead of a volumetric ray-trace, we flatten the 8-layer grid topology
        (8 strands, 8x8 grid) into a contiguous 1D audio-buffer style array for
        hardware-accelerated parallel MAC (Multiply-Accumulate) operations.
        """
        self.num_strands = num_strands
        self.grid_size = grid_size
        self.total_elements = num_strands * grid_size * grid_size
        
        # 1D flattened buffer representing the 8-layer Braid of Grids.
        # This acts exactly like a multi-channel digital audio buffer.
        self.state_buffer = np.zeros(self.total_elements, dtype=np.int32)
        
        # Sidon address set for each strand (fixed for 8-strand braid)
        self.sidon_addresses = [sidon_address(k) for k in range(num_strands)]
        
    def get_index(self, strand, x, y):
        return strand * (self.grid_size * self.grid_size) + y * self.grid_size + x
        
    def inject_signal(self, strand, x, y, energy_q16):
        """
        Injects energy into a specific node.
        """
        idx = self.get_index(strand, x, y)
        self.state_buffer[idx] += int(energy_q16 * Q16_SCALE)
        
    def compute_collisions_simd(self):
        """
        Executes a parallel DSP SIMD pass to compute topological collisions (Mountain Merges).
        Since strands are layered vertically (z-axis), a collision occurs when multiple
        strands hold energy in the same (x, y) coordinates.
        
        By treating the buffer as 8 interleaved channels, we can compute collisions
        by convolving a vertical "ray" across the flattened array, analogous to an FIR filter.
        """
        # Reshape to (8, 64) to allow vectorized column operations
        layered_view = self.state_buffer.reshape((self.num_strands, self.grid_size * self.grid_size))
        
        # Compute vertical energy accumulation using pure SIMD
        column_energy = np.sum(layered_view, axis=0)
        
        # Threshold MAC equivalent: Find where column energy exceeds a collision threshold
        collision_threshold = 2 * Q16_SCALE
        collisions = column_energy > collision_threshold
        
        collision_indices = np.where(collisions)[0]
        
        results = []
        for c_idx in collision_indices:
            x = c_idx % self.grid_size
            y = c_idx // self.grid_size
            results.append({
                "x": int(x),
                "y": int(y),
                "energy": int(column_energy[c_idx])
            })
            
        return results

    # ─── QR decomposition as VCN braid crossings ───────────────────────────

    def qr_factorize(self, A: list[list[float]]) -> dict:
        """
        QR decomposition of an 8×8 matrix using Householder reflectors.
        
        Each reflector is a braid crossing in the VCN pipeline:
          Householder H_k at column k → crossing at VCN strand k
          Q = H_0 · ... · H_7 → product braid word
          R = upper triangular → post-crossing residual
        
        Returns dict with Q, R, reflectors, and VCN strand mapping.
        """
        m = len(A)
        n = len(A[0])
        k_min = min(m, n)
        
        R = [row[:] for row in A]
        reflectors = []
        
        for col in range(k_min):
            x = [R[i][col] for i in range(col, m)]
            v_raw, tau = _householder_vector(x)
            if abs(tau) < EPSILON:
                continue
            
            v_full = [0.0] * col + v_raw
            reflectors.append(HouseholderReflector(k=col, v=v_full, tau=tau))
            
            for j in range(col, n):
                dot = sum(v_full[i] * R[i][j] for i in range(col, m))
                for i in range(col, m):
                    R[i][j] -= tau * v_full[i] * dot
        
        # Construct Q = product of Householder reflectors
        Q = [[float(i == j) for j in range(m)] for i in range(m)]
        for ref in reversed(reflectors):
            for j in range(m):
                dot = sum(ref.v[i] * Q[i][j] for i in range(m))
                for i in range(m):
                    Q[i][j] -= ref.tau * ref.v[i] * dot
        
        # Clean near-zeros
        for i in range(m):
            for j in range(n):
                if abs(R[i][j]) < EPSILON:
                    R[i][j] = 0.0
                if abs(Q[i][j]) < EPSILON:
                    Q[i][j] = 0.0
        
        # Map reflectors to VCN strands
        strand_map = []
        for ref in reflectors:
            strand_map.append({
                "vcn_strand": ref.k,
                "sidon_address": ref.sidon,
                "sidon_slack": sidon_address(7) - ref.sidon,
                "tau": round(ref.tau, 6),
                "nonzeros": sum(1 for vi in ref.v if abs(vi) > EPSILON),
            })
        
        # QR error
        A_reconstructed = np.dot(np.array(Q), np.array(R))
        error = np.linalg.norm(np.array(A) - A_reconstructed)
        
        return {
            "Q": [[round(v, 8) for v in row] for row in Q],
            "R": [[round(v, 8) for v in row] for row in R],
            "n_reflectors": len(reflectors),
            "qr_error": round(error, 12),
            "vcn_strand_mapping": strand_map,
            "total_sidon_weight": sum(ref.sidon for ref in reflectors),
            "sidon_utilization": f"{sum(ref.sidon for ref in reflectors)}/255",
        }

    def load_matrix_into_state(self, A: list[list[float]]):
        """Load an 8×8 matrix into the VCN state buffer as strand energies."""
        self.state_buffer = np.zeros(self.total_elements, dtype=np.int32)
        n = min(len(A), self.num_strands)
        for i in range(n):
            for j in range(min(len(A[i]), self.grid_size)):
                self.inject_signal(strand=i, x=j, y=0, energy_q16=A[i][j])

if __name__ == "__main__":
    print("=" * 60)
    print("VCN DSP SIMD Pipeline — QR Decomposition via Braid Crossings")
    print("=" * 60)
    
    # ─── Part 1: Original collision detection ─────────────────────────────
    pipeline = VCNDSPPipeline()
    
    print("\n[1/3] Topological collision detection...")
    pipeline.inject_signal(strand=0, x=3, y=4, energy_q16=1.5)
    pipeline.inject_signal(strand=2, x=3, y=4, energy_q16=1.0)
    pipeline.inject_signal(strand=1, x=7, y=7, energy_q16=0.5)
    pipeline.inject_signal(strand=7, x=7, y=7, energy_q16=0.5)
    pipeline.inject_signal(strand=3, x=1, y=1, energy_q16=2.5)
    
    collisions = pipeline.compute_collisions_simd()
    print(f"  Detected {len(collisions)} topological collisions")
    for c in collisions:
        print(f"    (x={c['x']}, y={c['y']}) energy={c['energy']/Q16_SCALE:.4f}")
    
    # ─── Part 2: QR decomposition as braid crossings ──────────────────────
    print("\n[2/3] QR decomposition via Householder braid crossings...")
    
    # Test matrices
    test_matrices = [
        ("I₈ (identity)", [[float(i==j) for j in range(8)] for i in range(8)]),
        ("diag(3,1,1,1,1,1,1,1)", [[3.0 if i==j and i==0 else (1.0 if i==j else 0.0) for j in range(8)] for i in range(8)]),
    ]
    
    import random
    random.seed(42)
    rand8 = [[random.uniform(-1, 1) for _ in range(8)] for _ in range(8)]
    test_matrices.append(("random 8×8", rand8))
    
    qr_results = []
    for name, A in test_matrices:
        result = pipeline.qr_factorize(A)
        qr_results.append({"name": name, **result})
        
        print(f"\n  Matrix: {name}")
        print(f"  Reflectors: {result['n_reflectors']} (braid crossings)")
        print(f"  QR error:   {result['qr_error']:.2e}")
        print(f"  Sidon util: {result['sidon_utilization']}")
        for s in result['vcn_strand_mapping']:
            print(f"    VCN strand {s['vcn_strand']:d}: sidon=2^{s['vcn_strand']}={s['sidon_address']:3d} "
                  f"τ={s['tau']:.4f} nz={s['nonzeros']}")
    
    # ─── Part 3: Load QR result into VCN state ───────────────────────────
    print("\n[3/3] Loading QR result into VCN state buffer...")
    for name, A in test_matrices[:1]:
        pipeline.load_matrix_into_state(A)
        # Verify by checking the loaded state
        layered = pipeline.state_buffer.reshape((pipeline.num_strands, pipeline.grid_size * pipeline.grid_size))
        loaded_energy = np.sum(layered, axis=1) / Q16_SCALE
        print(f"  {name} loaded — strand energies: {['{:.2f}'.format(e) for e in loaded_energy]}")
    
    # ─── Receipt ─────────────────────────────────────────────────────────
    receipt = {
        "schema": "vcn_qr_braid_v1",
        "claim_boundary": "householder_qr_on_vcn_pipeline;each_reflector_is_braid_crossing",
        "sidon_addresses": pipeline.sidon_addresses,
        "collision_results": collisions,
        "qr_results": qr_results,
        "summary": {
            "total_tests": len(test_matrices),
            "max_error": max(r["qr_error"] for r in qr_results),
            "braid_word_length": f"≤8 crossings per 8×8 decomposition",
            "vcn_strands": pipeline.num_strands,
            "vcn_grid": f"{pipeline.grid_size}×{pipeline.grid_size}",
        },
        "computed_at": datetime.now(timezone.utc).isoformat(),
    }
    canonical = json.dumps(receipt, sort_keys=True, separators=(",", ":"))
    receipt["receipt_sha256"] = hashlib.sha256(canonical.encode()).hexdigest()
    
    with open("vcn_qr_braid_receipt.json", "w") as f:
        json.dump(receipt, f, indent=2)
    
    print(f"\n{'='*60}")
    print(f"Receipt: vcn_qr_braid_receipt.json")
    print(f"SHA256:  {receipt['receipt_sha256']}")