#!/usr/bin/env python3
"""
Q4_K_M Tensor Regeneration from DNA Seed
Deterministic tensor-aware synthesis with receipts

Format: GGUF-compatible Q4_K_M blocks
- 32-value groups
- 4-bit quantization codes
- Per-block scale + zero
- Sign mask for negative values

MODES:
- LEGACY: Uses rng.integers() - produces procedural noise (NOT weights)
- BRAID:  Uses compiled folding maps - produces ACTUAL weight reconstruction
"""
import hashlib
import struct
import json
import argparse
from pathlib import Path
from datetime import datetime
import numpy as np

# Import BRAID compiler for structure-preserving regeneration
try:
    from helix_cdc.braid_compiler import BraidCompiler
    BRAID_AVAILABLE = True
except ImportError:
    BRAID_AVAILABLE = False

def load_centroid_atlas(atlas_path: Path) -> np.ndarray:
    """Load centroid atlas (16-bin magnitude centroids)"""
    if atlas_path.exists():
        return np.load(atlas_path)
    else:
        # Generate synthetic v0 atlas
        print(f"⚠️  Atlas not found, generating synthetic v0")
        return generate_synthetic_atlas_v0()

def generate_synthetic_atlas_v0() -> np.ndarray:
    """Generate synthetic 16-bin magnitude atlas (seeded)"""
    np.random.seed(0x1F647D91)  # From DNA seed

    # 16 magnitude bins, exponentially spaced
    centroids = np.array([
        0.0001, 0.001, 0.005, 0.01,
        0.02, 0.05, 0.1, 0.2,
        0.4, 0.8, 1.5, 3.0,
        6.0, 12.0, 24.0, 48.0
    ], dtype=np.float32)

    return centroids

def combine_seeds(dna_hex: str, entropy_seed: str) -> str:
    """Combine DNA + entropy seeds"""
    combined = hashlib.sha256(f"{entropy_seed}:{dna_hex}".encode()).hexdigest()
    return combined

def deterministic_rng_context(combined_seed: str, label: str, offset: int) -> np.random.Generator:
    """Create deterministic RNG from seed + label + offset (FibPi3D)"""
    context = f"{combined_seed}:{label}:{offset}"
    context_hash = hashlib.sha256(context.encode()).digest()

    # Use first 8 bytes as uint64 seed
    rng_seed = struct.unpack('<Q', context_hash[:8])[0]

    return np.random.default_rng(rng_seed)

def synthesize_q4km_block(rng: np.random.Generator, atlas: np.ndarray, block_size: int = 32) -> bytes:
    """
    Synthesize one Q4_K_M block (32 values) - OPTIMIZED FOR LOW MEMORY

    Format:
    - 2 bytes: scale (fp16)
    - 2 bytes: zero (fp16)
    - 4 bytes: sign mask (32 bits for 32 values)
    - 16 bytes: 4-bit codes (2 per byte, 32 codes)
    Total: 24 bytes per block
    """
    atlas_len = len(atlas)

    # Pre-allocate lists instead of numpy arrays
    magnitude_codes = []
    values = []
    sign_mask = 0

    # Generate values one at a time to minimize memory
    for i in range(block_size):
        code = rng.integers(0, atlas_len, dtype=np.uint8)
        magnitude = atlas[code]
        sign = rng.integers(0, 2, dtype=np.uint8)

        magnitude_codes.append(code)
        value = magnitude if sign == 0 else -magnitude
        values.append(value)

        if sign:
            sign_mask |= (1 << i)

    # Compute block scale and zero
    min_val = min(values)
    max_val = max(values)

    scale = (max_val - min_val) / 15.0 if max_val != min_val else 1.0
    zero = min_val

    # Pack block
    block = bytearray(24)  # Pre-allocate exact size

    # Scale + Zero (fp16)
    struct.pack_into('<e', block, 0, scale)
    struct.pack_into('<e', block, 2, zero)

    # Sign mask (32 bits)
    struct.pack_into('<I', block, 4, sign_mask)

    # 4-bit codes (2 per byte)
    for i in range(0, block_size, 2):
        byte = (magnitude_codes[i] & 0xF) | ((magnitude_codes[i+1] & 0xF) << 4)
        block[8 + i//2] = byte

    return bytes(block)

def grow_q4km_weights(
    dna_hex: str,
    entropy_seed: str,
    label: str,
    offset: int,
    size: int,
    atlas: np.ndarray,
    output_file
) -> dict:
    """
    Generate Q4_K_M tensor weights from DNA seed (STREAMING)

    Args:
        dna_hex: DNA seed (hex string, e.g., "1f647d91f6")
        entropy_seed: Entropy seed (e.g., "bf690b54")
        label: Tensor label (e.g., "blk.00.attn.q")
        offset: FibPi3D fold offset
        size: Total bytes to generate
        atlas: Centroid atlas
        output_file: File handle to write to (must support .write())

    Returns:
        receipt_dict
    """
    # Combine seeds
    combined_seed = combine_seeds(dna_hex, entropy_seed)

    # Create deterministic RNG
    rng = deterministic_rng_context(combined_seed, label, offset)

    # Generate Q4_K_M blocks (STREAMING)
    block_size_bytes = 24  # bytes per Q4_K_M block
    num_blocks = (size + block_size_bytes - 1) // block_size_bytes

    sha = hashlib.sha256()
    bytes_written = 0

    for i in range(num_blocks):
        block = synthesize_q4km_block(rng, atlas, block_size=32)

        # Write only what we need (trim last block if needed)
        if bytes_written + len(block) > size:
            block = block[:size - bytes_written]

        output_file.write(block)
        sha.update(block)
        bytes_written += len(block)

        # Flush periodically to avoid buffer bloat
        if i % 10000 == 0:
            output_file.flush()

    # Generate receipt
    receipt = {
        "version": "regrow_q4km_v0_streaming",
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "inputs": {
            "dna_hex": dna_hex,
            "entropy_seed": entropy_seed,
            "label": label,
            "offset": offset,
            "size": size
        },
        "atlas": {
            "type": "template-atlas synthetic" if atlas.shape[0] == 16 else "Crystal Centroids v1",
            "bins": int(atlas.shape[0]),
            "atlas_hash": hashlib.sha256(atlas.tobytes()).hexdigest()[:16]
        },
        "output": {
            "size": bytes_written,
            "blocks": num_blocks,
            "sha256": sha.hexdigest()
        },
        "determinism": {
            "combined_seed": combined_seed[:32] + "...",
            "rng_context": f"{combined_seed[:16]}:{label}:{offset}"
        },
        "format": {
            "type": "Q4_K_M",
            "block_size": 32,
            "block_bytes": 24,
            "gguf_compatible": "in spirit (container wiring next)"
        }
    }

    return receipt


# ========== FibPi3D PARALLEL REGEN (STEP 2 INTEGRATION) ==========

def _regen_chunk_worker(seed_ctx: dict, job: dict) -> bytes:
    """
    Worker function for FibPi scheduler: regenerate a single chunk.

    Args:
        seed_ctx: Context with dna_hex, entropy_seed, atlas, base_offset
        job: Job dict with name, offset, size

    Returns:
        bytes: Regenerated chunk

    NOTE: This is called by multiprocessing workers, so it must be:
    - Pure/deterministic (same inputs → same output)
    - Picklable (no lambdas, no local closures)

    CRITICAL: Uses base_offset for RNG seed (not chunk offset) to maintain determinism.
    Then skips to correct block position within the stream.
    """
    from io import BytesIO

    dna_hex = seed_ctx["dna_hex"]
    entropy_seed = seed_ctx["entropy_seed"]
    atlas = seed_ctx["atlas"]
    base_offset = seed_ctx["base_offset"]  # RNG seed offset (NOT chunk offset)

    label = job["name"]
    chunk_offset = job["offset"]  # Where this chunk starts in the full tensor
    size = job["size"]

    # CRITICAL: Use base_offset for RNG seed (determinism)
    combined_seed = combine_seeds(dna_hex, entropy_seed)
    rng = deterministic_rng_context(combined_seed, label, base_offset)

    # Calculate how many Q4_K_M blocks to skip to reach this chunk
    block_size_bytes = 24
    blocks_to_skip = chunk_offset // block_size_bytes
    blocks_to_generate = (size + block_size_bytes - 1) // block_size_bytes

    # Skip ahead in RNG stream (generate and discard blocks)
    for _ in range(blocks_to_skip):
        synthesize_q4km_block(rng, atlas, block_size=32)

    # Now generate the blocks we actually need
    buffer = BytesIO()
    bytes_written = 0

    for _ in range(blocks_to_generate):
        block = synthesize_q4km_block(rng, atlas, block_size=32)

        # Trim last block if needed
        if bytes_written + len(block) > size:
            block = block[:size - bytes_written]

        buffer.write(block)
        bytes_written += len(block)

    return buffer.getvalue()


def grow_q4km_weights_parallel(
    dna_hex: str,
    entropy_seed: str,
    label: str,
    offset: int,
    size: int,
    atlas: np.ndarray,
    output_file,
    chunk_size: int = 65520,  # 2730 blocks * 24 bytes = 65520 bytes (block-aligned!)
    use_fibpi: bool = True
) -> dict:
    """
    PARALLEL version of grow_q4km_weights using FibPi3D scheduler + UDR.

    Splits regen work into chunks, schedules via FibPi wave-front geometry,
    and executes in parallel across CPU cores (GPU/QPU future work).

    Args:
        (same as grow_q4km_weights, plus:)
        chunk_size: Chunk size for parallelization (default 64KB)
        use_fibpi: If True, use FibPi scheduler; if False, fall back to linear

    Returns:
        receipt_dict (same as grow_q4km_weights, plus scheduler metrics)

    Raises:
        RuntimeError: If scheduler not available or GPU/QPU requested
    """
    import sys
    import os
    import time

    # Add helix-cdc to path for scheduler
    helix_cdc_path = os.path.expanduser("~/helix-cdc")
    if helix_cdc_path not in sys.path:
        sys.path.insert(0, helix_cdc_path)

    try:
        from helix_cdc.scheduler.regen_fibpi_scheduler import (
            schedule_regen_jobs,
            create_chunked_jobs
        )
    except ImportError as e:
        raise RuntimeError(
            "FibPi scheduler not available. Install helix-cdc or use grow_q4km_weights() instead."
        ) from e

    t0 = time.time()

    # 1) Create chunked jobs
    seed_ctx = {
        "dna_hex": dna_hex,
        "entropy_seed": entropy_seed,
        "atlas": atlas,
        "base_offset": offset  # CRITICAL: Use original offset for RNG seed (determinism)
    }

    jobs = create_chunked_jobs(
        name=label,
        total_offset=offset,
        total_size=size,
        chunk_size=chunk_size,
        regen_fn=_regen_chunk_worker,
        seed_ctx=seed_ctx
    )

    print(f"   FibPi scheduler: {len(jobs)} chunks ({chunk_size:,} bytes each)")

    # 2) Check if C backend should be used (DEFAULT: enabled)
    use_c_backend = os.getenv("HELIX_USE_C_REGEN", "1") == "1"
    if use_c_backend:
        print(f"   Using C backend (helix_regen_cli) for regen (144 MB/s)")

    # 3) Schedule and execute
    topology = "fibpi" if use_fibpi else "linear"
    results = schedule_regen_jobs(
        jobs,
        topology=topology,
        N=4,
        seed=f"helix_regen:{label}",
        use_c_backend=use_c_backend
    )

    # 4) Write chunks in order
    sha = hashlib.sha256()
    bytes_written = 0

    for result in sorted(results, key=lambda r: r["offset"]):
        chunk_bytes = result["bytes"]
        output_file.write(chunk_bytes)
        sha.update(chunk_bytes)
        bytes_written += len(chunk_bytes)

    dt = time.time() - t0

    # 4) Generate receipt (augmented with scheduler info)
    combined_seed = combine_seeds(dna_hex, entropy_seed)

    receipt = {
        "version": "regrow_q4km_v0_fibpi_parallel",
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "inputs": {
            "dna_hex": dna_hex,
            "entropy_seed": entropy_seed,
            "label": label,
            "offset": offset,
            "size": size
        },
        "atlas": {
            "type": "template-atlas synthetic" if atlas.shape[0] == 16 else "Crystal Centroids v1",
            "bins": int(atlas.shape[0]),
            "atlas_hash": hashlib.sha256(atlas.tobytes()).hexdigest()[:16]
        },
        "output": {
            "size": bytes_written,
            "sha256": sha.hexdigest()
        },
        "scheduler": {
            "topology": topology,
            "num_chunks": len(jobs),
            "chunk_size": chunk_size,
            "backends": [r["backend"] for r in results],
            "Se_range": [min(r["Se"] for r in results), max(r["Se"] for r in results)],
            "duration_s": round(dt, 3),
            "chunks_per_sec": round(len(jobs) / dt, 2) if dt > 0 else 0,
            "implementation": results[0]["implementation"] if results else "unknown"
        },
        "determinism": {
            "combined_seed": combined_seed[:32] + "...",
            "rng_context": f"{combined_seed[:16]}:{label}:{offset}"
        },
        "format": {
            "type": "Q4_K_M",
            "block_size": 32,
            "block_bytes": 24,
            "gguf_compatible": "in spirit (container wiring next)"
        }
    }

    return receipt


# ========== BRAID-MODE REGENERATION (STRUCTURE, NOT NOISE) ==========

def compile_weights_to_braid_seed(
    weights: np.ndarray,
    tensor_name: str,
    codebook_size: int = 256
) -> dict:
    """
    Compile weights to a BRAID seed that contains reconstruction data.

    This is the INVERSE of regrow - it encodes actual weight structure,
    not random noise.

    Args:
        weights: Original tensor weights to encode
        tensor_name: Name of the tensor
        codebook_size: Size of quantization codebook

    Returns:
        BRAID seed dict with folding maps
    """
    if not BRAID_AVAILABLE:
        raise RuntimeError("BRAID compiler not available. Check helix_cdc.braid_compiler")

    compiler = BraidCompiler(codebook_size=codebook_size)
    seed = compiler.compile_seed(weights, tensor_name=tensor_name)

    return seed


def regrow_weights_from_braid_seed(seed: dict) -> np.ndarray:
    """
    Regrow weights from a BRAID seed using stored folding maps.

    This is DETERMINISTIC reconstruction from structure, not random generation.

    Args:
        seed: BRAID seed dict with folding maps

    Returns:
        Reconstructed weights
    """
    if not BRAID_AVAILABLE:
        raise RuntimeError("BRAID compiler not available. Check helix_cdc.braid_compiler")

    compiler = BraidCompiler()
    return compiler.regrow_from_seed(seed)


def verify_braid_regeneration(
    original: np.ndarray,
    seed: dict
) -> dict:
    """
    Verify that BRAID regeneration recovers original weights.

    Returns verification metrics.
    """
    if not BRAID_AVAILABLE:
        raise RuntimeError("BRAID compiler not available. Check helix_cdc.braid_compiler")

    compiler = BraidCompiler()
    return compiler.verify_reconstruction(original, seed)


def grow_q4km_weights_braid(
    seed: dict,
    output_file
) -> dict:
    """
    BRAID-mode: Generate Q4_K_M tensor weights from BRAID seed.

    This uses stored folding maps for DETERMINISTIC reconstruction,
    not random generation like the legacy mode.

    Args:
        seed: BRAID seed dict with ::FOLDING_MAP:: etc
        output_file: File handle to write reconstructed weights

    Returns:
        Receipt dict with regeneration metrics
    """
    if not BRAID_AVAILABLE:
        raise RuntimeError("BRAID compiler not available. Check helix_cdc.braid_compiler")

    import time
    t0 = time.time()

    # Regrow from seed
    compiler = BraidCompiler()
    weights = compiler.regrow_from_seed(seed)

    # Convert to bytes
    weight_bytes = weights.astype(np.float32).tobytes()

    # Write
    output_file.write(weight_bytes)

    dt = time.time() - t0

    # Generate receipt
    receipt = {
        "version": "regrow_q4km_v1_braid",
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "mode": "BRAID",
        "seed": {
            "tensor_name": seed.get("::TENSOR_NAME::", "unknown"),
            "cosmic_anchor": seed.get("::COSMIC_ANCHOR::", "")[:16] + "...",
            "weight_hash": seed.get("::WEIGHT_HASH::", "")[:16] + "...",
            "codebook_size": len(seed["::ANALYSIS::"]["codebook"]),
        },
        "output": {
            "size": len(weight_bytes),
            "shape": list(weights.shape),
            "sha256": hashlib.sha256(weight_bytes).hexdigest(),
        },
        "reconstruction": {
            "method": "BRAID folding map",
            "residual_count": seed["::FOLDING_MAP::"]["residual_count"],
            "duration_s": round(dt, 3),
        },
        "format": {
            "type": "float32_raw",
            "note": "Can be quantized to Q4_K_M if needed",
        }
    }

    return receipt


def main():
    parser = argparse.ArgumentParser(description="Generate Q4_K_M tensor from DNA seed")
    parser.add_argument("--dna", required=True, help="DNA seed (hex)")
    parser.add_argument("--entropy", required=True, help="Entropy seed (hex)")
    parser.add_argument("--label", required=True, help="Tensor label")
    parser.add_argument("--offset", type=int, default=0, help="FibPi3D offset")
    parser.add_argument("--size", type=int, required=True, help="Size in bytes")
    parser.add_argument("--atlas", type=Path, default=Path("centroid_atlas_v0.npy"), help="Atlas path")
    parser.add_argument("--out", type=Path, required=True, help="Output file")
    parser.add_argument("--fibpi", action="store_true", default=True,
                        help="Use FibPi3D wave scheduling (DEFAULT: enabled)")
    parser.add_argument("--no-fibpi", action="store_true",
                        help="Disable FibPi3D, use single-threaded streaming")
    parser.add_argument("--chunk-size", type=int, default=65520,
                        help="Chunk size for FibPi parallel regen (default: 65520)")

    args = parser.parse_args()

    # --no-fibpi overrides --fibpi
    use_fibpi = args.fibpi and not args.no_fibpi

    print("=" * 80)
    print("Q4_K_M TENSOR REGENERATION")
    print("=" * 80)
    print()

    # Load atlas
    print(f"📚 Loading centroid atlas...")
    atlas = load_centroid_atlas(args.atlas)
    print(f"   Atlas: {atlas.shape[0]} bins")
    print(f"   Range: {atlas.min():.4f} to {atlas.max():.4f}")
    print()

    # Generate tensor
    print(f"🧬 Generating tensor...")
    print(f"   DNA: {args.dna}")
    print(f"   Entropy: {args.entropy}")
    print(f"   Label: {args.label}")
    print(f"   Offset: {args.offset}")
    print(f"   Size: {args.size:,} bytes")
    print(f"   Mode: {'FibPi3D WAVE (multi-core)' if use_fibpi else 'STREAMING (single-core)'}")
    print()

    print(f"💾 {'Wave regen' if use_fibpi else 'Streaming'} to {args.out}...")

    if use_fibpi:
        # FibPi3D WAVE SCHEDULING - uses all cores
        with open(args.out, 'wb') as f:
            receipt = grow_q4km_weights_parallel(
                dna_hex=args.dna,
                entropy_seed=args.entropy,
                label=args.label,
                offset=args.offset,
                size=args.size,
                atlas=atlas,
                output_file=f,
                chunk_size=args.chunk_size,
                use_fibpi=True
            )
    else:
        # Single-threaded streaming (legacy)
        with open(args.out, 'wb') as f:
            receipt = grow_q4km_weights(
                dna_hex=args.dna,
                entropy_seed=args.entropy,
                label=args.label,
                offset=args.offset,
                size=args.size,
                atlas=atlas,
                output_file=f
            )

    print(f"   ✅ Complete")
    print(f"   Size: {receipt['output']['size']:,} bytes")
    print(f"   SHA256: {receipt['output']['sha256']}")

    # Save receipt
    receipt_path = args.out.with_suffix(args.out.suffix + '.receipt.json')
    receipt_path.write_text(json.dumps(receipt, indent=2))
    print(f"   Receipt: {receipt_path}")
    print()

    print("=" * 80)
    print("✅ COMPLETE")
    print("=" * 80)
    print()
    print("Verify:")
    print(f"  sha256sum {args.out}")
    print(f"  cat {receipt_path}")
    print()

if __name__ == "__main__":
    main()
