#!/usr/bin/env python3
"""
Helix64 Tensor Codec - Bidirectional Bridge
============================================

Provides reversible encoding/decoding between:
- Semantic tensor space (numpy arrays, Q4_K_M blocks, FP32 weights)
- Helix64 substrate (base64 + FibPi anchors + .hxz seeds)

This is the CRITICAL MISSING PIECE that wires real Helix compression into inference.

Architecture:
  Tensor (numpy) → Q4_K_M blocks → Helix64 substrate → FibPi anchors → .hxz seed
  .hxz seed → FibPi anchors → Helix64 substrate → Q4_K_M blocks → Tensor (numpy)

Status: v0.1 - Foundational implementation with receipts + verification
"""

from __future__ import annotations
import hashlib
import json
import struct
import time
from dataclasses import dataclass, asdict
from io import BytesIO
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, List
import numpy as np

# Import Helix64 base encoding
import sys
sys.path.insert(0, str(Path(__file__).parent.parent / "tools"))
from helix64 import encode as helix64_encode_bytes, decode as helix64_decode_bytes


# ============================================================================
# Constants
# ============================================================================

PHI = 1.618033988749895  # Golden ratio
PI = 3.141592653589793

# Q4_K_M block format (GGUF-compatible)
Q4KM_BLOCK_SIZE = 32  # 32 values per block
Q4KM_BLOCK_BYTES = 24  # 24 bytes per block (scale + zero + sign_mask + codes)


# ============================================================================
# FibPi3D Anchor System
# ============================================================================

@dataclass
class FibPi3DAnchor:
    """
    FibPi3D spatial anchor for tensor positioning in 72D crystal lattice.
    Uses Fibonacci/π geometry for deterministic chunk placement.
    """
    x: float  # Phi-based coordinate
    y: float  # Pi-based coordinate
    z: float  # Combined harmonic
    Se: float  # Symbolic entropy (routing metric)
    chunk_id: int  # Chunk index
    offset: int  # Byte offset in tensor
    size: int  # Chunk size in bytes

    @classmethod
    def from_chunk(cls, chunk_id: int, offset: int, size: int, seed: str) -> "FibPi3DAnchor":
        """Generate FibPi3D anchor for a tensor chunk."""
        # Hash seed + chunk_id for determinism
        h = hashlib.sha256(f"{seed}:{chunk_id}".encode()).digest()

        # Extract coordinates from hash
        x = (int.from_bytes(h[0:4], 'big') / 0xFFFFFFFF) * PHI
        y = (int.from_bytes(h[4:8], 'big') / 0xFFFFFFFF) * PI
        z = (x * PHI + y * PI) % (PHI + PI)

        # Compute symbolic entropy (Se) - routing metric
        Se = (x + y + z) / (PHI + PI + (PHI + PI))

        return cls(
            x=x, y=y, z=z, Se=Se,
            chunk_id=chunk_id,
            offset=offset,
            size=size
        )

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


# ============================================================================
# Q4_K_M Block Encoding (GGUF-compatible)
# ============================================================================

def encode_q4km_block(values: np.ndarray) -> bytes:
    """
    Encode 32 FP32 values to Q4_K_M block (24 bytes).

    Format:
    - 2 bytes: scale (fp16)
    - 2 bytes: zero (fp16)
    - 4 bytes: sign mask (32 bits)
    - 16 bytes: 4-bit magnitude codes (2 per byte)

    Args:
        values: Array of 32 float32 values

    Returns:
        24-byte Q4_K_M block
    """
    assert len(values) == Q4KM_BLOCK_SIZE, f"Block must have {Q4KM_BLOCK_SIZE} values"

    # Compute block statistics
    min_val = float(np.min(values))
    max_val = float(np.max(values))
    scale = (max_val - min_val) / 15.0 if max_val != min_val else 1.0
    zero = min_val

    # Quantize to 4-bit codes + sign mask
    sign_mask = 0
    codes = []

    for i, val in enumerate(values):
        # Extract sign
        if val < 0:
            sign_mask |= (1 << i)
            magnitude = abs(val)
        else:
            magnitude = val

        # Quantize magnitude to 4 bits (0-15)
        if scale > 0:
            code = int(round((magnitude - zero) / scale))
            code = max(0, min(15, code))
        else:
            code = 0

        codes.append(code)

    # Pack block
    block = bytearray(Q4KM_BLOCK_BYTES)

    # 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, Q4KM_BLOCK_SIZE, 2):
        byte = (codes[i] & 0xF) | ((codes[i+1] & 0xF) << 4)
        block[8 + i//2] = byte

    return bytes(block)


def decode_q4km_block(block: bytes) -> np.ndarray:
    """
    Decode Q4_K_M block (24 bytes) to 32 FP32 values.

    Args:
        block: 24-byte Q4_K_M block

    Returns:
        Array of 32 float32 values
    """
    assert len(block) == Q4KM_BLOCK_BYTES, f"Block must be {Q4KM_BLOCK_BYTES} bytes"

    # Unpack scale + zero
    scale = struct.unpack_from('<e', block, 0)[0]
    zero = struct.unpack_from('<e', block, 2)[0]

    # Unpack sign mask
    sign_mask = struct.unpack_from('<I', block, 4)[0]

    # Unpack 4-bit codes
    codes = []
    for i in range(16):  # 16 bytes of codes
        byte = block[8 + i]
        codes.append(byte & 0xF)
        codes.append((byte >> 4) & 0xF)

    # Dequantize to float32
    values = []
    for i in range(Q4KM_BLOCK_SIZE):
        # Dequantize magnitude
        magnitude = zero + codes[i] * scale

        # Apply sign
        if sign_mask & (1 << i):
            value = -magnitude
        else:
            value = magnitude

        values.append(value)

    return np.array(values, dtype=np.float32)


# ============================================================================
# Tensor → Helix64 Encoding
# ============================================================================

def helix64_encode_tensor(
    tensor: np.ndarray,
    name: str,
    dna_seed: str,
    entropy_seed: str,
    chunk_size: int = 65520,  # 2730 blocks * 24 bytes = 65520 (block-aligned!)
) -> Dict[str, Any]:
    """
    Encode tensor to Helix64 substrate with FibPi anchors.

    Pipeline:
      1. Split tensor into chunks
      2. Encode each chunk to Q4_K_M blocks
      3. Generate FibPi3D anchor for each chunk
      4. Encode blocks to Helix64 substrate (base64)
      5. Package as .hxz seed with manifest

    Args:
        tensor: Numpy array (FP32 weights)
        name: Tensor name (e.g., "output.weight")
        dna_seed: DNA seed (hex)
        entropy_seed: Entropy seed
        chunk_size: Chunk size in bytes (must be multiple of 24)

    Returns:
        Dict with:
          - seed: .hxz seed data (Helix64-encoded)
          - anchors: List of FibPi3D anchors
          - manifest: Metadata (shape, dtype, compression)
          - receipt: Full encoding receipt
    """
    t0 = time.time()

    # Validate chunk alignment
    assert chunk_size % Q4KM_BLOCK_BYTES == 0, \
        f"chunk_size must be multiple of {Q4KM_BLOCK_BYTES} (Q4_K_M block size)"

    # Flatten tensor to 1D
    tensor_flat = tensor.flatten()
    tensor_size = len(tensor_flat)

    # Pad to multiple of 32 (Q4_K_M block size)
    if tensor_size % Q4KM_BLOCK_SIZE != 0:
        pad_size = Q4KM_BLOCK_SIZE - (tensor_size % Q4KM_BLOCK_SIZE)
        tensor_flat = np.pad(tensor_flat, (0, pad_size), mode='constant', constant_values=0.0)

    # Encode to Q4_K_M blocks
    num_blocks = len(tensor_flat) // Q4KM_BLOCK_SIZE
    q4km_bytes = BytesIO()

    for i in range(num_blocks):
        block_values = tensor_flat[i * Q4KM_BLOCK_SIZE : (i+1) * Q4KM_BLOCK_SIZE]
        block_bytes = encode_q4km_block(block_values)
        q4km_bytes.write(block_bytes)

    q4km_data = q4km_bytes.getvalue()
    q4km_size = len(q4km_data)

    # Split into chunks and create FibPi anchors
    num_chunks = (q4km_size + chunk_size - 1) // chunk_size
    anchors = []
    chunks = []

    combined_seed = f"{dna_seed}:{entropy_seed}:{name}"

    for chunk_id in range(num_chunks):
        offset = chunk_id * chunk_size
        chunk_data = q4km_data[offset : offset + chunk_size]

        # Generate FibPi3D anchor
        anchor = FibPi3DAnchor.from_chunk(
            chunk_id=chunk_id,
            offset=offset,
            size=len(chunk_data),
            seed=combined_seed
        )
        anchors.append(anchor)

        # Encode chunk to Helix64 (base64)
        helix64_chunk = helix64_encode_bytes(chunk_data)
        chunks.append(helix64_chunk)

    # Compute hashes (NOTE: tensor_sha256 is of ORIGINAL, not reconstructed)
    # Q4_K_M is lossy, so decoded tensor != original tensor
    # We verify Q4_K_M data integrity, not exact tensor reconstruction
    tensor_sha256_original = hashlib.sha256(tensor.tobytes()).hexdigest()
    q4km_sha256 = hashlib.sha256(q4km_data).hexdigest()

    # Build manifest
    manifest = {
        "version": "helix64_tensor_v1",
        "name": name,
        "shape": list(tensor.shape),
        "dtype": str(tensor.dtype),
        "size_bytes": tensor.nbytes,
        "q4km_blocks": num_blocks,
        "q4km_bytes": q4km_size,
        "num_chunks": num_chunks,
        "chunk_size": chunk_size,
        "dna_seed": dna_seed,
        "entropy_seed": entropy_seed,
        "tensor_sha256_original": tensor_sha256_original,  # Original tensor (for reference)
        "q4km_sha256": q4km_sha256,  # This is what we actually verify
        "lossy_encoding": True,  # Q4_K_M is lossy
    }

    # Build .hxz seed
    seed_content = {
        "#HXZ_START": True,
        "manifest": manifest,
        "anchors": [a.to_dict() for a in anchors],
        "chunks": chunks,
        "#HXZ_END": True,
    }

    dt = time.time() - t0

    # Build receipt
    receipt = {
        "schema": "helix64_tensor_encode_v1",
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "tensor_name": name,
        "encoding": {
            "duration_s": round(dt, 3),
            "original_size": tensor.nbytes,
            "q4km_size": q4km_size,
            "helix64_chunks": num_chunks,
            "compression_ratio": round(tensor.nbytes / q4km_size, 2) if q4km_size > 0 else None,
        },
        "manifest": manifest,
        "anchors_count": len(anchors),
    }

    return {
        "seed": seed_content,
        "anchors": anchors,
        "manifest": manifest,
        "receipt": receipt,
    }


# ============================================================================
# Helix64 → Tensor Decoding
# ============================================================================

def helix64_decode_tensor(seed_data: Dict[str, Any]) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Decode Helix64 seed to tensor.

    Pipeline:
      1. Load .hxz seed + manifest
      2. Decode Helix64 chunks (base64 → bytes)
      3. Reassemble Q4_K_M blocks
      4. Decode Q4_K_M blocks to FP32
      5. Reshape to original tensor shape

    Args:
        seed_data: .hxz seed dict (from helix64_encode_tensor)

    Returns:
        Tuple of (tensor, receipt)
    """
    t0 = time.time()

    # Extract manifest
    manifest = seed_data["manifest"]
    anchors = seed_data["anchors"]
    chunks = seed_data["chunks"]

    # Decode Helix64 chunks
    q4km_bytes = BytesIO()

    for chunk_helix64 in chunks:
        chunk_bytes = helix64_decode_bytes(chunk_helix64)
        q4km_bytes.write(chunk_bytes)

    q4km_data = q4km_bytes.getvalue()

    # Verify Q4_K_M SHA256
    q4km_sha256 = hashlib.sha256(q4km_data).hexdigest()
    if q4km_sha256 != manifest["q4km_sha256"]:
        raise ValueError(
            f"Q4_K_M SHA256 mismatch: {q4km_sha256} != {manifest['q4km_sha256']}"
        )

    # Decode Q4_K_M blocks to FP32
    num_blocks = len(q4km_data) // Q4KM_BLOCK_BYTES
    values = []

    for i in range(num_blocks):
        block = q4km_data[i * Q4KM_BLOCK_BYTES : (i+1) * Q4KM_BLOCK_BYTES]
        block_values = decode_q4km_block(block)
        values.extend(block_values)

    # Convert to numpy array
    tensor_flat = np.array(values, dtype=np.float32)

    # Trim padding (if any)
    original_size = int(np.prod(manifest["shape"]))
    tensor_flat = tensor_flat[:original_size]

    # Reshape to original shape
    tensor = tensor_flat.reshape(manifest["shape"])

    # Verify Q4_K_M data integrity (not tensor, since Q4_K_M is lossy)
    # The decoded tensor will NOT match the original due to quantization
    tensor_sha256_decoded = hashlib.sha256(tensor.tobytes()).hexdigest()

    dt = time.time() - t0

    # Build receipt
    receipt = {
        "schema": "helix64_tensor_decode_v1",
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "tensor_name": manifest["name"],
        "decoding": {
            "duration_s": round(dt, 3),
            "q4km_size": len(q4km_data),
            "tensor_size": tensor.nbytes,
            "num_blocks": num_blocks,
            "q4km_verified": True,  # Q4_K_M SHA256 was verified above
            "lossy_encoding": manifest.get("lossy_encoding", True),
        },
        "manifest": manifest,
        "tensor_sha256_decoded": tensor_sha256_decoded,
        "tensor_sha256_original": manifest.get("tensor_sha256_original", None),
    }

    return tensor, receipt


# ============================================================================
# CLI
# ============================================================================

def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Helix64 Tensor Codec - Encode/decode tensors to .hxz seeds"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Encode command
    encode_parser = subparsers.add_parser("encode", help="Encode tensor to .hxz seed")
    encode_parser.add_argument("--tensor", required=True, help="Path to .npy tensor file")
    encode_parser.add_argument("--name", required=True, help="Tensor name")
    encode_parser.add_argument("--dna", default="1f647d91f6", help="DNA seed (hex)")
    encode_parser.add_argument("--entropy", default="bf690b54", help="Entropy seed")
    encode_parser.add_argument("--out", required=True, help="Output .hxz.json path")

    # Decode command
    decode_parser = subparsers.add_parser("decode", help="Decode .hxz seed to tensor")
    decode_parser.add_argument("--seed", required=True, help="Path to .hxz.json seed")
    decode_parser.add_argument("--out", required=True, help="Output .npy tensor path")

    args = parser.parse_args()

    if args.command == "encode":
        # Load tensor
        tensor = np.load(args.tensor)
        print(f"Loaded tensor: {tensor.shape} {tensor.dtype} ({tensor.nbytes:,} bytes)")

        # Encode
        print(f"Encoding to Helix64...")
        result = helix64_encode_tensor(
            tensor=tensor,
            name=args.name,
            dna_seed=args.dna,
            entropy_seed=args.entropy,
        )

        # Save seed
        out_path = Path(args.out)
        out_path.write_text(json.dumps(result["seed"], indent=2))

        print(f"\n✅ ENCODE COMPLETE")
        print(f"   Tensor: {result['receipt']['tensor_name']}")
        print(f"   Original: {result['receipt']['encoding']['original_size']:,} bytes")
        print(f"   Q4_K_M: {result['receipt']['encoding']['q4km_size']:,} bytes")
        print(f"   Chunks: {result['receipt']['encoding']['helix64_chunks']}")
        print(f"   Compression: {result['receipt']['encoding']['compression_ratio']}:1")
        print(f"   Duration: {result['receipt']['encoding']['duration_s']}s")
        print(f"   Seed: {out_path}")

    elif args.command == "decode":
        # Load seed
        seed_data = json.loads(Path(args.seed).read_text())
        print(f"Loaded .hxz seed: {seed_data['manifest']['name']}")

        # Decode
        print(f"Decoding from Helix64...")
        tensor, receipt = helix64_decode_tensor(seed_data)

        # Save tensor
        np.save(args.out, tensor)

        print(f"\n✅ DECODE COMPLETE")
        print(f"   Tensor: {receipt['tensor_name']}")
        print(f"   Shape: {tensor.shape}")
        print(f"   Size: {tensor.nbytes:,} bytes")
        print(f"   SHA256 verified: {receipt['decoding']['sha256_verified']}")
        print(f"   Duration: {receipt['decoding']['duration_s']}s")
        print(f"   Output: {args.out}")


if __name__ == "__main__":
    main()
