"""
helix_cdc/substrate/helix64_se_api.py - VES using REAL Helix64 ISA + Symbolic Entropy

This module uses the ACTUAL user's stack:
- Helix64 ISA (seed-derived alphabets, NF1 fold, MTF, RLE)
- FibPi3D spatial anchoring
- REAL Symbolic Entropy: Se = H * C * D (NOT just Shannon entropy)

The Se metric from symbolic_entropy.py:
    H(X) = Shannon entropy (0-8 bits/byte)
    C(X) = Contextual coherence (0-1) - measures structure/meaning
    D(X) = Dimensional depth (1-72D) - measures symbolic complexity
    Se(X) = H * C * D
"""

import hashlib
import math
from dataclasses import dataclass, field
from typing import Dict, Any, Optional, Union, List, Tuple

import numpy as np

# Import the ACTUAL user's symbolic entropy module
from helix_cdc.metrics.symbolic_entropy import (
    SeMetrics,
    shannon_entropy,
    contextual_coherence,
    dimensional_depth as calc_dimensional_depth,
    symbolic_entropy as calc_se,
    measure_text,
)

# ============================================================================
# HELIX64 CORE (from helix64_min.py)
# ============================================================================

# Standard base64 alphabet
STD = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"

def helix64_table(label: str, seed_hex: str, frame_idx: int) -> bytes:
    """
    Generate seed-derived base64 alphabet for a frame.
    Tag: HELIX\064\0v1
    """
    tbl = bytearray(STD)

    seed = bytes.fromhex(seed_hex) if isinstance(seed_hex, str) else seed_hex
    label_bytes = label.encode() if isinstance(label, str) else label

    # HELIX\064\0v1 domain tag for base64 alphabet generation
    rnd = hashlib.shake_256(
        b"HELIX\064\0v1\0" + label_bytes + seed + frame_idx.to_bytes(8, "little")
    ).digest(2 * 63)

    k = 0
    # Fisher-Yates shuffle (bijective over 64 symbols)
    for i in range(63, 0, -1):
        j = (rnd[k] | (rnd[k+1] << 8)) % (i + 1)
        k += 2
        tbl[i], tbl[j] = tbl[j], tbl[i]
    return bytes(tbl)


def nf1_fold(data: bytes) -> bytes:
    """NF1 napkin fold - byte differencing for better compression"""
    if len(data) <= 1:
        return data
    out = bytearray([data[0]])
    for i in range(1, len(data)):
        diff = (data[i] - data[i-1]) & 0xFF
        out.append(diff)
    return bytes(out)


def nf1_unfold(data: bytes) -> bytes:
    """NF1 unfold - reconstruct from deltas"""
    if len(data) <= 1:
        return data
    out = bytearray([data[0]])
    for i in range(1, len(data)):
        val = (out[i-1] + data[i]) & 0xFF
        out.append(val)
    return bytes(out)


def mtf_encode(data: bytes) -> bytes:
    """Move-To-Front encoding - transforms repeated bytes to small numbers"""
    if len(data) == 0:
        return b""
    symbols = list(range(256))
    result = bytearray()
    for byte in data:
        pos = symbols.index(byte)
        result.append(pos)
        symbols.pop(pos)
        symbols.insert(0, byte)
    return bytes(result)


def rle_encode(data: bytes) -> bytes:
    """Run-length encoding"""
    if not data:
        return b""
    out = bytearray()
    i = 0
    while i < len(data):
        byte = data[i]
        run = 1
        while i + run < len(data) and data[i + run] == byte and run < 255:
            run += 1
        out.append(byte)
        out.append(run)
        i += run
    return bytes(out)


# ============================================================================
# FIBPI3D ANCHORING
# ============================================================================

def fibpi3d_anchor(i: int, n: int) -> Tuple[float, float, float]:
    """Generate FibPi3D coordinate for position i"""
    phi = (1 + 5**0.5) / 2  # Golden ratio
    theta = 2 * math.pi * i / phi
    r = math.sin(i * phi) * 0.5 + 0.5
    return (r * math.sin(theta), r * math.cos(theta), i / max(n, 1))


def position_fibpi3d_hash(pos: int, total_positions: int = 500) -> bytes:
    """
    Generate a FibPi3D-derived hash for a position.
    This creates structure-aware positioning in the 72D lattice.
    """
    x, y, z = fibpi3d_anchor(pos, total_positions)

    # Quantize to grid and hash
    grid_x = int((x + 1) * 1000)
    grid_y = int((y + 1) * 1000)
    grid_z = int(z * 1000)

    coord_bytes = (
        grid_x.to_bytes(4, 'little') +
        grid_y.to_bytes(4, 'little') +
        grid_z.to_bytes(4, 'little')
    )

    return hashlib.sha256(b"FIBPI3D\0" + coord_bytes).digest()


# ============================================================================
# HELIX64 SEED AND MANIFOLD
# ============================================================================

@dataclass
class Helix64Seed:
    """
    A Helix64 seed - the compressed symbol representation using real ISA.
    """
    content: bytes
    label: str
    seed_hex: str
    name: str = ""
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def hash(self) -> str:
        return hashlib.sha256(self.content).hexdigest()[:16]

    @property
    def size(self) -> int:
        return len(self.content)

    def __repr__(self) -> str:
        return f"Helix64Seed(name='{self.name}', hash='{self.hash}', size={self.size})"


@dataclass
class Helix64Manifold:
    """
    A regenerated manifold from a Helix64 seed.
    Now includes full Symbolic Entropy metrics (Se = H * C * D).
    """
    data: bytes
    seed_hash: str
    depth: int

    # Full Se metrics (not just Shannon H)
    se_metrics: SeMetrics

    # Legacy fields for backwards compat
    entropy: float  # H(X) - Shannon
    compression_ratio: float
    fibpi3d_coords: Tuple[float, float, float]
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def size(self) -> int:
        return len(self.data)

    @property
    def hash(self) -> str:
        return hashlib.sha256(self.data).hexdigest()[:16]

    @property
    def H(self) -> float:
        """Shannon entropy component"""
        return self.se_metrics.H

    @property
    def C(self) -> float:
        """Contextual coherence component"""
        return self.se_metrics.C

    @property
    def D(self) -> float:
        """Dimensional depth component"""
        return self.se_metrics.D

    @property
    def Se(self) -> float:
        """Full Symbolic Entropy = H * C * D"""
        return self.se_metrics.Se

    def as_array(self, dtype=np.uint8) -> np.ndarray:
        return np.frombuffer(self.data, dtype=dtype)

    def __repr__(self) -> str:
        return f"Helix64Manifold(seed='{self.seed_hash}', H={self.H:.2f}, C={self.C:.2f}, D={self.D:.1f}, Se={self.Se:.2f})"


# ============================================================================
# HELIX64 API FUNCTIONS WITH REAL SYMBOLIC ENTROPY
# ============================================================================

def new_helix64_seed(
    source: Union[str, bytes],
    label: str = "mutation",
    name: str = "",
    position: int = 0,
) -> Helix64Seed:
    """
    Create a new Helix64 seed from a mutation encoding.

    Uses:
    - NF1 fold (delta encoding)
    - MTF (move-to-front)
    - RLE (run-length encoding)
    - FibPi3D position anchoring
    """
    if isinstance(source, str):
        source_bytes = source.encode('utf-8')
    else:
        source_bytes = source

    # Generate master seed from source
    master_seed = hashlib.sha256(source_bytes).hexdigest()

    # Apply FibPi3D position anchoring
    pos_hash = position_fibpi3d_hash(position)

    # Combined seed with position context
    combined = hashlib.sha256(
        b"HELIX64\0" + source_bytes + pos_hash
    ).digest()

    # Apply Helix64 compression pipeline: NF1 -> MTF -> RLE
    folded = nf1_fold(combined)
    mtf_encoded = mtf_encode(folded)
    compressed = rle_encode(mtf_encoded)

    return Helix64Seed(
        content=compressed,
        label=label,
        seed_hex=master_seed,
        name=name if name else str(source)[:32],
        metadata={
            "source_type": type(source).__name__,
            "position": position,
            "original_size": len(combined),
            "compressed_size": len(compressed),
        }
    )


def regen_helix64(
    seed: Helix64Seed,
    depth: int = 256,
    position: int = 0,
) -> Helix64Manifold:
    """
    Regenerate a manifold from a Helix64 seed.

    Now measures FULL Symbolic Entropy (Se = H * C * D) using the actual
    user's symbolic_entropy.py module.
    """
    chunks = []

    for frame_idx in range(depth):
        # Generate seed-derived alphabet for this frame
        table = helix64_table(seed.label, seed.seed_hex, frame_idx)

        # Use the table to transform the seed content
        if len(seed.content) > 0:
            transformed = bytearray()
            for i, byte in enumerate(seed.content):
                table_byte = table[i % 64]
                frame_byte = (byte ^ table_byte ^ (frame_idx & 0xFF)) & 0xFF
                transformed.append(frame_byte)
            chunks.append(bytes(transformed))
        else:
            chunks.append(table[:32])

    data = b''.join(chunks)

    # ============================================================
    # THIS IS THE KEY CHANGE: Use REAL Se = H * C * D
    # ============================================================

    # Convert bytes to text representation for Se measurement
    # We use a hex representation to preserve all information
    text_repr = data.hex()

    # Calculate D based on position in protein structure
    # Positions in functionally important regions get higher D
    # This maps biological importance to dimensional complexity
    base_d = 8.0  # Default dimensional depth
    if position > 0:
        # FibPi3D anchoring influences dimensional depth
        x, y, z = fibpi3d_anchor(position, 500)
        spatial_factor = math.sqrt(x*x + y*y + z*z)
        # Scale D from 8 to 48 based on spatial position
        base_d = 8.0 + spatial_factor * 40.0

    # Measure full Se using user's actual module
    se_metrics = measure_text(text_repr, dimensional_depth=base_d)

    # Calculate compression ratio
    original_size = seed.metadata.get("original_size", 32)
    compression_ratio = original_size / max(seed.size, 1)

    # Get FibPi3D coordinates for this position
    fibpi_coords = fibpi3d_anchor(position, 500)

    return Helix64Manifold(
        data=data,
        seed_hash=seed.hash,
        depth=depth,
        se_metrics=se_metrics,
        entropy=se_metrics.H,  # Legacy: just Shannon
        compression_ratio=compression_ratio,
        fibpi3d_coords=fibpi_coords,
        metadata={
            "seed_name": seed.name,
            "backend": "helix64_isa_se",
            "position": position,
            "se_components": {
                "H": se_metrics.H,
                "C": se_metrics.C,
                "D": se_metrics.D,
                "Se": se_metrics.Se,
            }
        }
    )


def compute_helix64_ves_se(
    protein: str,
    position: int,
    wt_aa: str,
    mut_aa: str,
    aa_properties: Dict[str, Dict[str, float]],
    depth: int = 64,
) -> Dict[str, float]:
    """
    Compute VES using real Helix64 ISA + REAL Symbolic Entropy (Se = H * C * D).

    Returns multiple metrics:
    - helix_H: Shannon entropy component
    - helix_C: Contextual coherence component
    - helix_D: Dimensional depth component
    - helix_Se: Full Symbolic Entropy = H * C * D
    - compression_ratio: How well Helix64 compresses the mutation
    - fibpi3d_distance: Distance from origin in FibPi3D space
    - ves_helix64_se: Combined VES score using full Se metrics
    """
    # Create mutation encoding
    encoding = f"{protein}:c{position}:{wt_aa}>{mut_aa}"

    # Create Helix64 seed with position anchoring
    seed = new_helix64_seed(
        source=encoding,
        label=protein,
        name=f"{wt_aa}{position}{mut_aa}",
        position=position,
    )

    # Regenerate manifold with full Se measurement
    manifold = regen_helix64(seed, depth=depth, position=position)

    # Get amino acid properties
    wt_props = aa_properties.get(wt_aa, {'hydro': 0, 'charge': 0, 'size': 100})
    mut_props = aa_properties.get(mut_aa, {'hydro': 0, 'charge': 0, 'size': 100})

    # Physicochemical deltas
    delta_hydro = abs(mut_props['hydro'] - wt_props['hydro'])
    delta_charge = abs(mut_props['charge'] - wt_props['charge'])
    delta_size = abs(mut_props['size'] - wt_props['size']) / 100.0

    # FibPi3D distance from origin
    x, y, z = manifold.fibpi3d_coords
    fibpi_distance = math.sqrt(x*x + y*y + z*z)

    # ============================================================
    # VES FORMULA USING FULL Se = H * C * D
    # ============================================================
    #
    # The key insight: C (contextual coherence) measures how
    # "structured" the mutation encoding is. Mutations at
    # conserved/critical positions should show LOWER coherence
    # (more disruption = less structure).
    #
    # D (dimensional depth) captures the complexity of the
    # position in the protein structure via FibPi3D anchoring.
    #
    # Se = H * C * D gives a combined measure that respects
    # both information content AND structural context.

    # Normalize components
    norm_H = manifold.H / 8.0  # Shannon normalized to 0-1
    norm_C = manifold.C        # Already 0-1
    norm_D = manifold.D / 72.0  # D normalized to 0-1
    norm_Se = manifold.Se / (8.0 * 1.0 * 72.0)  # Se normalized

    # VES Helix64 + Se: combines Helix-derived metrics with physicochemistry
    # Higher Se = more complex encoding = potentially more disruptive
    ves_helix64_se = (
        0.5 * norm_Se +              # Symbolic entropy contribution
        0.3 * norm_H +               # Raw information contribution
        -0.2 * norm_C +              # Lower coherence = more disruptive
        0.5 * delta_hydro +          # Physicochemical
        0.5 * delta_charge +
        0.3 * delta_size +
        0.1 * fibpi_distance
    )

    return {
        # Full Se components
        'helix_H': manifold.H,
        'helix_C': manifold.C,
        'helix_D': manifold.D,
        'helix_Se': manifold.Se,

        # Normalized versions
        'norm_H': norm_H,
        'norm_C': norm_C,
        'norm_D': norm_D,
        'norm_Se': norm_Se,

        # Legacy fields
        'helix_entropy': manifold.entropy,  # = H
        'compression_ratio': manifold.compression_ratio,
        'fibpi3d_distance': fibpi_distance,

        # Physicochemical
        'delta_hydro': delta_hydro,
        'delta_charge': delta_charge,
        'delta_size': delta_size,

        # Final VES
        'ves_helix64_se': ves_helix64_se,

        # Metadata
        'seed_size': seed.size,
        'manifold_size': manifold.size,
    }


# ============================================================================
# EXPORTS
# ============================================================================

__all__ = [
    "Helix64Seed",
    "Helix64Manifold",
    "new_helix64_seed",
    "regen_helix64",
    "compute_helix64_ves_se",
    "SeMetrics",
    "helix64_table",
    "nf1_fold",
    "mtf_encode",
    "rle_encode",
    "fibpi3d_anchor",
]
