#!/usr/bin/env python3
"""
Helix64 Phi-Codebook: Seeds as Regeneration Instructions
=========================================================

This codec aligns with the CDR mental model:
- Seeds ARE the instructions, not pointers to a lookup table
- Basis patterns are GENERATED from φ-lattice, not learned
- Selection uses Se (symbolic entropy) steering, not nearest-centroid
- No residuals - the substrate does the work

Architecture:
1. Generate crystal basis from FibPi3D coordinates (deterministic)
2. Select SPARSE subset of basis patterns using greedy + Se penalty
3. Store only: which seeds + coefficients (tiny)
4. Decode: regenerate basis from seeds, apply coefficients

This is the "Einstein layer" over the Newtonian VQ approach.

Usage:
    codec = Helix64PhiCodebook(num_bases=64, sparsity=8)
    encoded = codec.encode(tensor)
    decoded = codec.decode(encoded)
"""

import math
import hashlib
import numpy as np
from typing import Dict, List, Tuple, Any, Optional
from dataclasses import dataclass

# Import existing substrate primitives
try:
    from helix_cdc.crystal_codec import (
        generate_crystal_basis,
        fibpi3d_coord,
        generate_basis_pattern,
        PHI, PI, GOLDEN_ANGLE
    )
    from helix_cdc.substrate.fibpi import (
        fibpi_mix_u64,
        fibpi_energy,
        fibpi_score
    )
    SUBSTRATE_AVAILABLE = True
except ImportError:
    SUBSTRATE_AVAILABLE = False
    PHI = 1.618033988749895
    PI = math.pi
    GOLDEN_ANGLE = 2.39996322972865332


# ============================================================================
# Se (Symbolic Entropy) Computation
# ============================================================================

def compute_se(values: np.ndarray, window: int = None) -> float:
    """
    Compute Symbolic Entropy Se = log1p(H × C × D)

    Where:
        H = Shannon entropy of value distribution
        C = Coherence (1 - normalized variance)
        D = Dimensional depth (based on spread of values)

    This is the control surface for seed selection.
    """
    if len(values) == 0:
        return 0.0

    # Normalize to [0, 1]
    v = values.flatten().astype(np.float64)
    v_min, v_max = v.min(), v.max()
    if v_max - v_min < 1e-10:
        return 0.0
    v_norm = (v - v_min) / (v_max - v_min + 1e-10)

    # H: Shannon entropy via histogram
    hist, _ = np.histogram(v_norm, bins=32, range=(0, 1))
    hist = hist / (hist.sum() + 1e-10)
    hist = hist[hist > 0]
    H = -np.sum(hist * np.log2(hist + 1e-10)) / 5.0  # Normalize to ~[0,1]
    H = np.clip(H, 0.0, 1.0)

    # C: Coherence (1 - normalized variance)
    var = np.var(v_norm)
    C = 1.0 - np.clip(var * 4, 0.0, 1.0)  # Scale variance

    # D: Dimensional depth (spread)
    D = np.clip(np.std(v_norm) * 3, 0.1, 1.0)

    # Se = log1p(H × C × D)
    Se = float(np.log1p(H * C * D))

    return Se


def se_penalty(basis_coeffs: np.ndarray, target_se_band: Tuple[float, float] = (0.3, 0.7)) -> float:
    """
    Compute penalty for Se outside target band.

    Lower penalty = better alignment with healthy Se range.
    This is the steering mechanism for basis selection.
    """
    se = compute_se(basis_coeffs)
    se_low, se_high = target_se_band

    if se_low <= se <= se_high:
        return 0.0  # In band - no penalty
    elif se < se_low:
        return (se_low - se) ** 2
    else:
        return (se - se_high) ** 2


# ============================================================================
# Phi-Codebook: Generative Basis from Seeds
# ============================================================================

@dataclass
class PhiEncoding:
    """Phi-codebook encoding of a tensor."""
    shape: Tuple[int, ...]
    seed: str
    num_bases: int
    sparsity: int  # How many bases actually used
    selected_indices: np.ndarray  # Which basis patterns selected
    coefficients: np.ndarray  # Coefficients for selected bases
    se_at_encode: float
    reconstruction_error: float

    def to_dict(self) -> Dict[str, Any]:
        return {
            "shape": list(self.shape),
            "seed": self.seed,
            "num_bases": self.num_bases,
            "sparsity": self.sparsity,
            "selected_indices": self.selected_indices.tolist(),
            "coefficients": self.coefficients.tolist(),
            "se_at_encode": self.se_at_encode,
            "reconstruction_error": self.reconstruction_error,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "PhiEncoding":
        return cls(
            shape=tuple(d["shape"]),
            seed=d["seed"],
            num_bases=d["num_bases"],
            sparsity=d["sparsity"],
            selected_indices=np.array(d["selected_indices"], dtype=np.int32),
            coefficients=np.array(d["coefficients"], dtype=np.float32),
            se_at_encode=d.get("se_at_encode", 0.0),
            reconstruction_error=d.get("reconstruction_error", 0.0),
        )

    def encoded_bytes(self) -> int:
        """Estimate encoded size in bytes."""
        # indices (int16) + coefficients (float32) + metadata
        return len(self.selected_indices) * 2 + len(self.coefficients) * 4 + 64


class Helix64PhiCodebook:
    """
    Phi-Codebook Codec: Seeds as Regeneration Instructions

    Key differences from VQ:
    - Basis patterns are GENERATED from FibPi3D, not learned from data
    - Selection uses Se steering, not nearest centroid
    - Sparse: only K bases used, not all N
    - No residuals in "substrate" mode
    """

    def __init__(
        self,
        num_bases: int = 64,
        sparsity: int = 8,  # How many bases to select
        seed: str = "phi_codebook",
        se_weight: float = 0.1,  # Weight for Se penalty in selection
        target_se_band: Tuple[float, float] = (0.3, 0.7),
    ):
        self.num_bases = num_bases
        self.sparsity = min(sparsity, num_bases)
        self.seed = seed
        self.se_weight = se_weight
        self.target_se_band = target_se_band

        # Cache for basis patterns
        self._basis_cache: Dict[int, np.ndarray] = {}

    def _get_basis(self, size: int) -> np.ndarray:
        """Get or generate crystal basis for given size."""
        if size in self._basis_cache:
            return self._basis_cache[size]

        if SUBSTRATE_AVAILABLE:
            basis = generate_crystal_basis(self.num_bases, size, self.seed)
        else:
            # Fallback: generate locally
            basis = self._generate_basis_fallback(size)

        self._basis_cache[size] = basis
        return basis

    def _generate_basis_fallback(self, size: int) -> np.ndarray:
        """Fallback basis generation without substrate imports."""
        patterns = []
        pattern_types = ["sinusoid", "polynomial", "wavelet", "fourier", "random"]

        for i in range(self.num_bases):
            # FibPi3D coordinate
            t = i * PHI
            angle = (t * GOLDEN_ANGLE) % (2 * PI)
            z = (i % 10000) / 10000.0
            r = math.sqrt(1 - z * z) if z < 1 else 0
            x, y = r * math.cos(angle), r * math.sin(angle)

            # Generate pattern
            ptype = pattern_types[i % len(pattern_types)]
            pattern = self._generate_pattern(x, y, z, size, ptype)
            patterns.append(pattern)

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

    def _generate_pattern(self, x: float, y: float, z: float, size: int, ptype: str) -> np.ndarray:
        """Generate a single basis pattern."""
        t = np.linspace(0, 1, size)

        if ptype == "sinusoid":
            freq_x = 1 + abs(x) * 10
            freq_y = 1 + abs(y) * 10
            phase = z * 2 * PI
            pattern = np.sin(2 * PI * freq_x * t + phase) * 0.5 + np.cos(2 * PI * freq_y * t) * 0.5
        elif ptype == "polynomial":
            pattern = x * t**2 + y * t + z
        elif ptype == "wavelet":
            sigma = 0.1 + abs(z) * 0.2
            center = 0.5 + x * 0.3
            pattern = np.exp(-((t - center)**2) / (2 * sigma**2)) * np.cos(2 * PI * (1 + abs(y) * 5) * t)
        elif ptype == "fourier":
            k = int(1 + abs(x * y * z) * 20) % 50 + 1
            pattern = np.cos(2 * PI * k * t + z * PI)
        else:
            seed_val = int(abs(x * 1000) + abs(y * 1000) + abs(z * 1000)) % (2**31)
            rng = np.random.RandomState(seed_val)
            pattern = rng.randn(size)

        norm = np.linalg.norm(pattern)
        if norm > 1e-8:
            pattern = pattern / norm

        return pattern.astype(np.float32)

    def _greedy_select_with_se(
        self,
        tensor_flat: np.ndarray,
        basis: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Greedy basis selection with Se penalty.

        This is the key algorithmic difference from VQ:
        - Select bases that minimize reconstruction error PLUS Se penalty
        - Prefer bases that keep coefficients in healthy Se band
        """
        n = len(tensor_flat)
        residual = tensor_flat.copy()

        selected = []
        coefficients = []
        used_indices = set()

        for _ in range(self.sparsity):
            best_idx = -1
            best_score = float('inf')
            best_coeff = 0.0

            for i in range(self.num_bases):
                if i in used_indices:
                    continue

                # Project residual onto this basis
                coeff = np.dot(residual, basis[i])

                # Reconstruction with this basis added
                new_residual = residual - coeff * basis[i]
                recon_error = np.mean(new_residual ** 2)

                # Se penalty on accumulated coefficients
                test_coeffs = np.array(coefficients + [coeff])
                se_pen = se_penalty(test_coeffs, self.target_se_band)

                # Combined score: reconstruction + Se steering
                score = recon_error + self.se_weight * se_pen

                if score < best_score:
                    best_score = score
                    best_idx = i
                    best_coeff = coeff

            if best_idx < 0:
                break

            # Accept this basis
            selected.append(best_idx)
            coefficients.append(best_coeff)
            used_indices.add(best_idx)

            # Update residual
            residual = residual - best_coeff * basis[best_idx]

        return np.array(selected, dtype=np.int32), np.array(coefficients, dtype=np.float32)

    def encode(self, tensor: np.ndarray, tensor_name: str = None) -> PhiEncoding:
        """
        Encode tensor using sparse basis selection with Se steering.

        Args:
            tensor: Input tensor (any shape)
            tensor_name: Optional name for debugging

        Returns:
            PhiEncoding with selected indices and coefficients
        """
        original_shape = tensor.shape
        flat = tensor.flatten().astype(np.float32)

        # Get basis for this size
        basis = self._get_basis(len(flat))

        # Se-steered greedy selection
        selected_indices, coefficients = self._greedy_select_with_se(flat, basis)

        # Compute reconstruction for metrics
        reconstructed = np.zeros_like(flat)
        for idx, coeff in zip(selected_indices, coefficients):
            reconstructed += coeff * basis[idx]

        recon_error = float(np.sqrt(np.mean((flat - reconstructed) ** 2)))
        se_at_encode = compute_se(coefficients)

        return PhiEncoding(
            shape=original_shape,
            seed=self.seed,
            num_bases=self.num_bases,
            sparsity=len(selected_indices),
            selected_indices=selected_indices,
            coefficients=coefficients,
            se_at_encode=se_at_encode,
            reconstruction_error=recon_error,
        )

    def decode(self, encoding: PhiEncoding) -> np.ndarray:
        """
        Decode from PhiEncoding back to tensor.

        Key: Basis patterns are REGENERATED from seed, not looked up.
        """
        flat_size = int(np.prod(encoding.shape))

        # Regenerate basis (deterministic from seed)
        basis = self._get_basis(flat_size)

        # Reconstruct from sparse selection
        reconstructed = np.zeros(flat_size, dtype=np.float32)
        for idx, coeff in zip(encoding.selected_indices, encoding.coefficients):
            reconstructed += coeff * basis[idx]

        return reconstructed.reshape(encoding.shape)

    def encode_decode(self, tensor: np.ndarray) -> Tuple[PhiEncoding, np.ndarray]:
        """Convenience: encode then decode."""
        enc = self.encode(tensor)
        dec = self.decode(enc)
        return enc, dec

    def compression_ratio(self, encoding: PhiEncoding) -> float:
        """Compute compression ratio."""
        original_bytes = int(np.prod(encoding.shape)) * 4  # FP32
        encoded_bytes = encoding.encoded_bytes()
        return original_bytes / encoded_bytes if encoded_bytes > 0 else 0.0


# ============================================================================
# Comparison Test
# ============================================================================

if __name__ == "__main__":
    import time

    print("=" * 70)
    print("Helix64 Phi-Codebook: Seeds as Regeneration Instructions")
    print("=" * 70)

    # Test tensor
    np.random.seed(42)
    size = 4096
    t = np.linspace(0, 1, size)

    # Structured signal (what neural weights look like)
    test_tensor = (
        np.sin(2 * PI * 3 * t) * 0.01 +
        np.cos(2 * PI * 7 * t) * 0.005 +
        np.sin(2 * PI * 11 * t + 0.5) * 0.003 +
        np.random.randn(size) * 0.002
    ).astype(np.float32)

    print(f"\nTest tensor: shape={test_tensor.shape}, bytes={test_tensor.nbytes}")
    print(f"Original Se: {compute_se(test_tensor):.4f}")

    # Test different sparsity levels
    for sparsity in [4, 8, 16, 32]:
        codec = Helix64PhiCodebook(num_bases=64, sparsity=sparsity, se_weight=0.1)

        t0 = time.time()
        encoding = codec.encode(test_tensor)
        t_encode = time.time() - t0

        t1 = time.time()
        decoded = codec.decode(encoding)
        t_decode = time.time() - t1

        # Metrics
        mse = np.mean((test_tensor - decoded) ** 2)
        rmse = np.sqrt(mse)
        rel_err = np.linalg.norm(test_tensor - decoded) / (np.linalg.norm(test_tensor) + 1e-10)
        compression = codec.compression_ratio(encoding)

        print(f"\nSparsity={sparsity}:")
        print(f"  Compression: {compression:.1f}:1")
        print(f"  RMSE: {rmse:.6f}")
        print(f"  Rel error: {rel_err:.4%}")
        print(f"  Se at encode: {encoding.se_at_encode:.4f}")
        print(f"  Selected bases: {list(encoding.selected_indices)}")
        print(f"  Encode time: {t_encode*1000:.1f}ms, Decode time: {t_decode*1000:.1f}ms")

    print("\n" + "=" * 70)
    print("This is CDR-style encoding: seeds ARE the regen instructions")
    print("No residuals, no lookup table - just sparse basis selection")
    print("=" * 70)
