#!/usr/bin/env python3
"""
BRAID Bridge for FibPi3D Scheduler
===================================

Wires BRAID compiler's folding maps into the FibPi3D scheduler.

This is THE MISSING LINK:
- FibPi3D handles WHEN (scheduling) and WHERE (72D position)
- BRAID handles WHAT (actual structure via folding maps)

Usage with scheduler:
    from helix_cdc.regrow.braid_bridge import (
        braid_regen_callback,
        compile_tensor_to_braid_seed,
        create_braid_jobs
    )

    # Compile original weights to BRAID seed (one-time)
    braid_seed = compile_tensor_to_braid_seed(original_weights, "layer_name")

    # Create jobs for scheduler
    jobs = create_braid_jobs(
        name="layer_name",
        total_size=len(original_weights.tobytes()),
        chunk_size=65536,
        braid_seed=braid_seed
    )

    # Schedule with FibPi3D
    results = schedule_regen_jobs(jobs, topology="fibpi")
"""

import numpy as np
import base64
import hashlib
from typing import Dict, Any, List, Optional, Tuple

from helix_cdc.braid_compiler import BraidCompiler


def braid_regen_callback(seed_ctx: Dict[str, Any], job: Dict[str, Any]) -> bytes:
    """
    BRAID-based regen callback for FibPi3D scheduler.

    This replaces the PRNG-based callbacks with structure-preserving regeneration.

    Args:
        seed_ctx: Contains:
            - braid_seed: The BRAID seed with folding maps
            - dtype: numpy dtype (default float32)
        job: Contains:
            - name: Tensor name
            - offset: Byte offset into tensor
            - size: Number of bytes to regenerate

    Returns:
        Regenerated bytes (ACTUAL structure, not noise)
    """
    braid_seed = seed_ctx.get("braid_seed")
    if braid_seed is None:
        raise ValueError("braid_seed required in seed_ctx for BRAID regen")

    dtype = seed_ctx.get("dtype", np.float32)

    # Get folding map data
    analysis = braid_seed.get("::ANALYSIS::")
    folding_map = braid_seed.get("::FOLDING_MAP::")

    if analysis is None or folding_map is None:
        raise ValueError("Invalid BRAID seed: missing ::ANALYSIS:: or ::FOLDING_MAP::")

    # Decode indices from base64
    indices = np.frombuffer(
        base64.b64decode(folding_map["indices_b64"]),
        dtype=np.uint8
    )

    # Get codebook
    codebook = np.array(analysis["codebook"], dtype=np.float32)

    # Calculate element offset and count
    element_size = np.dtype(dtype).itemsize
    start_element = job["offset"] // element_size
    num_elements = job["size"] // element_size

    # Bounds check
    if start_element + num_elements > len(indices):
        # Pad with last valid index if needed
        valid_end = min(start_element + num_elements, len(indices))
        indices_slice = indices[start_element:valid_end]
        if len(indices_slice) < num_elements:
            padding = np.full(num_elements - len(indices_slice), indices[-1], dtype=np.uint8)
            indices_slice = np.concatenate([indices_slice, padding])
    else:
        indices_slice = indices[start_element:start_element + num_elements]

    # Regrow from codebook lookup (NOT random!)
    values = codebook[indices_slice]

    # Apply residuals if available
    sig_residuals = folding_map.get("significant_residuals", {})
    for pos_str, residual in sig_residuals.items():
        pos = int(pos_str)
        local_pos = pos - start_element
        if 0 <= local_pos < len(values):
            values[local_pos] += residual

    # Convert to bytes
    return values.astype(dtype).tobytes()


def compile_tensor_to_braid_seed(
    weights: np.ndarray,
    tensor_name: str,
    codebook_size: int = 256
) -> Dict[str, Any]:
    """
    Compile a tensor to BRAID seed format.

    Args:
        weights: numpy array of weights
        tensor_name: name of the tensor
        codebook_size: size of codebook (default 256)

    Returns:
        BRAID seed dict with folding maps
    """
    compiler = BraidCompiler(codebook_size=codebook_size)
    return compiler.compile_seed(weights, tensor_name)


def create_braid_jobs(
    name: str,
    total_size: int,
    chunk_size: int,
    braid_seed: Dict[str, Any],
    dtype: np.dtype = np.float32,
    total_offset: int = 0
) -> List[Dict[str, Any]]:
    """
    Create BRAID-based jobs for FibPi3D scheduler.

    Args:
        name: Tensor name
        total_size: Total bytes to regenerate
        chunk_size: Chunk size per job
        braid_seed: BRAID seed with folding maps
        dtype: numpy dtype
        total_offset: Starting offset

    Returns:
        List of jobs ready for schedule_regen_jobs()
    """
    num_chunks = (total_size + chunk_size - 1) // chunk_size
    jobs = []

    seed_ctx = {
        "braid_seed": braid_seed,
        "dtype": dtype
    }

    for i in range(num_chunks):
        chunk_offset = total_offset + i * chunk_size
        chunk_end = min(chunk_offset + chunk_size, total_offset + total_size)
        chunk_len = chunk_end - chunk_offset

        jobs.append({
            "name": name,
            "offset": chunk_offset,
            "size": chunk_len,
            "fn": braid_regen_callback,
            "seed_ctx": seed_ctx
        })

    return jobs


def verify_braid_regen(
    original: np.ndarray,
    braid_seed: Dict[str, Any],
    chunk_size: int = 65536
) -> Dict[str, Any]:
    """
    Verify BRAID regeneration against original.

    Args:
        original: Original weights
        braid_seed: BRAID seed
        chunk_size: Chunk size for regen

    Returns:
        Verification results
    """
    jobs = create_braid_jobs(
        name="verify",
        total_size=original.nbytes,
        chunk_size=chunk_size,
        braid_seed=braid_seed,
        dtype=original.dtype
    )

    # Execute jobs sequentially (no scheduler for verification)
    regen_chunks = []
    for job in jobs:
        chunk = braid_regen_callback(job["seed_ctx"], job)
        regen_chunks.append(chunk)

    # Reassemble
    regen_bytes = b"".join(regen_chunks)
    regen_array = np.frombuffer(regen_bytes, dtype=original.dtype).reshape(original.shape)

    # Compare
    diff = original - regen_array
    mse = float(np.mean(diff ** 2))
    max_error = float(np.max(np.abs(diff)))
    rel_error = np.sqrt(mse) / (np.std(original) + 1e-10)

    orig_hash = hashlib.sha256(original.tobytes()).hexdigest()
    regen_hash = hashlib.sha256(regen_bytes).hexdigest()

    return {
        "mse": mse,
        "max_error": max_error,
        "relative_error": rel_error,
        "relative_error_pct": f"{rel_error * 100:.4f}%",
        "hashes_match": orig_hash == regen_hash,
        "original_hash": orig_hash,
        "regen_hash": regen_hash,
        "num_chunks": len(jobs),
        "chunk_size": chunk_size
    }


def demo():
    """Demo BRAID bridge with FibPi3D scheduler."""
    print("=" * 60)
    print("BRAID BRIDGE DEMO: Wiring Folding Maps to FibPi3D")
    print("=" * 60)

    # Create test weights
    np.random.seed(42)
    weights = np.random.randn(64, 64).astype(np.float32) * 0.1

    print(f"\n1. Original weights: {weights.shape}, {weights.nbytes} bytes")

    # Compile to BRAID seed
    print("\n2. Compiling to BRAID seed...")
    braid_seed = compile_tensor_to_braid_seed(weights, "test_layer")
    print(f"   Codebook size: {len(braid_seed['::ANALYSIS::']['codebook'])}")
    print(f"   Folding map entries: {len(base64.b64decode(braid_seed['::FOLDING_MAP::']['indices_b64']))}")

    # Create jobs
    print("\n3. Creating BRAID jobs for scheduler...")
    jobs = create_braid_jobs(
        name="test_layer",
        total_size=weights.nbytes,
        chunk_size=1024,  # 1KB chunks
        braid_seed=braid_seed
    )
    print(f"   Created {len(jobs)} jobs")

    # Verify
    print("\n4. Verifying chunked regeneration...")
    verification = verify_braid_regen(weights, braid_seed, chunk_size=1024)

    print(f"   MSE: {verification['mse']:.2e}")
    print(f"   Relative error: {verification['relative_error_pct']}")
    print(f"   Hashes match: {verification['hashes_match']}")

    # Show comparison with PRNG approach
    print("\n5. COMPARISON TO PRNG (what legacy code does):")
    rng = np.random.default_rng(42)
    prng_weights = rng.standard_normal((64, 64)).astype(np.float32) * 0.1
    prng_mse = float(np.mean((weights - prng_weights) ** 2))
    print(f"   PRNG MSE: {prng_mse:.2e}")
    print(f"   BRAID MSE: {verification['mse']:.2e}")
    print(f"   Improvement: {prng_mse / (verification['mse'] + 1e-20):.0f}x better")

    print("\n" + "=" * 60)
    print("BRAID BRIDGE READY FOR FIBPI3D SCHEDULER")
    print("=" * 60)

    return verification


if __name__ == "__main__":
    demo()
