#!/usr/bin/env python3
"""
codec_router.py - Se-Style Routing for Codec Chunks

Routes compression jobs based on chunk characteristics, similar to how
SeController routes prompts. Computes cheap metrics → chooses a lane.

Metrics (H/C/D for chunks):
    H = entropy/roughness (high for int8-like noise)
    C = structure/coherence (autocorrelation, low-freq energy)
    D = complexity proxy (spectral flatness, kurtosis)

Routes to:
    - Backend: gpu_batched_qr vs cpu
    - Basis: dct vs fibpi3d vs hybrid
    - Parameters: num_bases, overlap, patch_k
    - Escape hatch: store_raw=True

Decision policy:
    cos >= 0.98: accept (no patch)
    0.90 <= cos < 0.98: small patch (patch_k=32)
    0.75 <= cos < 0.90: more bases + patch
    cos < 0.75: escape hatch (store raw)
"""

from __future__ import annotations

import numpy as np
from dataclasses import dataclass, field
from typing import Optional, Literal, List, Dict, Any
from enum import Enum
import time


class BasisFamily(Enum):
    DCT = "dct"
    FIBPI3D = "fibpi3d"
    HYBRID = "hybrid"
    RAW = "raw"  # Escape hatch


class Backend(Enum):
    GPU_BATCHED = "gpu_batched_qr"
    CPU = "cpu"


@dataclass
class ChunkStats:
    """Cheap-to-compute metrics for a chunk (H/C/D style)."""

    # H - Entropy/roughness
    entropy: float = 0.0           # Approximate Shannon entropy
    value_range: float = 0.0       # max - min

    # C - Structure/coherence
    autocorr_lag1: float = 0.0     # First-order autocorrelation
    low_freq_ratio: float = 0.0   # Energy in bottom 25% of spectrum

    # D - Complexity proxy
    spectral_flatness: float = 0.0  # Geometric/arithmetic mean of spectrum
    kurtosis: float = 0.0          # Tailedness (high = spiky)
    sign_changes: float = 0.0      # Rate of sign changes

    # Size info
    size: int = 0

    @classmethod
    def compute(cls, chunk: np.ndarray) -> "ChunkStats":
        """Compute all stats for a chunk (designed to be fast)."""
        chunk = chunk.astype(np.float32)
        n = len(chunk)

        # H - Entropy approximation (histogram-based)
        hist, _ = np.histogram(chunk, bins=256)
        hist = hist[hist > 0]
        probs = hist / hist.sum()
        entropy = -np.sum(probs * np.log2(probs + 1e-10))

        # Value range
        value_range = float(chunk.max() - chunk.min())

        # C - Autocorrelation at lag 1
        if n > 1:
            mean = chunk.mean()
            var = chunk.var()
            if var > 1e-10:
                autocorr = np.corrcoef(chunk[:-1], chunk[1:])[0, 1]
            else:
                autocorr = 0.0
        else:
            autocorr = 0.0

        # Low frequency energy ratio (via FFT)
        fft = np.abs(np.fft.rfft(chunk))
        total_energy = np.sum(fft ** 2)
        low_freq_idx = len(fft) // 4
        low_freq_energy = np.sum(fft[:low_freq_idx] ** 2)
        low_freq_ratio = low_freq_energy / (total_energy + 1e-10)

        # D - Spectral flatness
        log_fft = np.log(fft + 1e-10)
        geometric_mean = np.exp(np.mean(log_fft))
        arithmetic_mean = np.mean(fft)
        spectral_flatness = geometric_mean / (arithmetic_mean + 1e-10)

        # Kurtosis
        std = chunk.std()
        if std > 1e-10:
            kurtosis = np.mean(((chunk - chunk.mean()) / std) ** 4) - 3
        else:
            kurtosis = 0.0

        # Sign change rate
        signs = np.sign(chunk)
        sign_changes = np.sum(signs[1:] != signs[:-1]) / (n - 1) if n > 1 else 0.0

        return cls(
            entropy=float(entropy),
            value_range=value_range,
            autocorr_lag1=float(autocorr) if not np.isnan(autocorr) else 0.0,
            low_freq_ratio=float(low_freq_ratio),
            spectral_flatness=float(spectral_flatness),
            kurtosis=float(kurtosis),
            sign_changes=float(sign_changes),
            size=n,
        )

    def is_structured(self) -> bool:
        """Heuristic: is this chunk structured (vs noise-like)?"""
        # High autocorrelation + high low-freq ratio = structured
        return self.autocorr_lag1 > 0.3 and self.low_freq_ratio > 0.5

    def is_noisy(self) -> bool:
        """Heuristic: is this chunk noise-like?"""
        # High entropy + low autocorrelation + high spectral flatness
        return (self.entropy > 6.0 and
                abs(self.autocorr_lag1) < 0.2 and
                self.spectral_flatness > 0.5)


@dataclass
class Budget:
    """Constraints for codec decisions.

    Modes:
        inference: Decode speed + behavioral fidelity (DEFAULT)
                   - Minimal decode ops (DCT inverse + add residual)
                   - Lossless reconstruction required
                   - Encode can be slow (one-time cost)

        archive: Maximum compression, accept quality loss
                 - Lossy, no exact residual
                 - Smaller storage footprint

        hybrid: Balance of speed and compression
                - Base + exact residual
                - Good compression with lossless decode
    """
    target_ms: float = 3000       # Target latency
    max_rss_mb: int = 8000        # Max RSS memory
    max_gpu_mb: int = 6000        # Max GPU memory
    min_cos: float = 0.99         # Minimum cosine similarity (stricter for inference)
    mode: Literal["inference", "archive", "hybrid"] = "inference"  # Default: decode speed + fidelity

    # Decode-speed tuning
    prefer_contiguous_residual: bool = True  # Avoid sparse lookups at decode time
    max_decode_ops: int = 2                   # Target: DCT inverse + add (2 ops)


class Lane(Enum):
    """Codec lanes (strategies)."""
    FAST_GPU_QR = "fast_gpu_qr"           # Batched QR, cached basis
    CPU_LSTSQ = "cpu_lstsq"               # Fallback CPU
    DCT_BASE = "dct_base"                 # DCT only
    DCT_BASE_PATCH = "dct_base+patch"     # DCT + sparse residual
    HYBRID_BASE_PATCH = "hybrid+patch"    # FibPi3D + DCT residual
    BASE_EXACT_RESIDUAL = "base+exact"    # Base + full residual (lossless)
    ESCAPE_RAW = "escape_raw"             # Store raw chunk


@dataclass
class CodecDecision:
    """What the router decided for a chunk."""

    lane: Lane = Lane.DCT_BASE
    backend: Backend = Backend.GPU_BATCHED
    basis: BasisFamily = BasisFamily.DCT
    num_bases: int = 256
    chunk_size: int = 4096
    overlap: int = 0
    patch_k: int = 0          # 0 = no patch
    store_exact_residual: bool = False  # For hybrid mode
    store_raw: bool = False   # Escape hatch

    # Quality targets
    target_cosine: float = 0.95

    # Diagnostics
    reason: str = ""
    stats: Optional[ChunkStats] = None
    budget: Optional[Budget] = None


@dataclass
class CodecResult:
    """Result of encoding a chunk."""

    coefficients: np.ndarray
    patch_indices: Optional[np.ndarray] = None
    patch_values: Optional[np.ndarray] = None
    raw_chunk: Optional[np.ndarray] = None  # If escape hatch used

    cosine_similarity: float = 0.0
    decision: Optional[CodecDecision] = None
    encode_time_ms: float = 0.0


class CodecRouter:
    """
    Se-style router for codec jobs.

    Routes chunks to appropriate backend/basis/parameters based on
    quick statistical analysis and budget constraints.

    Modes:
        - inference: Must run like original (near-lossless)
        - archive: Smaller storage, accept quality loss
        - hybrid: Store base + exact residual on demand
    """

    def __init__(
        self,
        default_budget: Optional[Budget] = None,
    ):
        self.default_budget = default_budget or Budget()

        # Cache for basis matrices
        self._basis_cache: Dict[tuple, np.ndarray] = {}
        self._qr_cache: Dict[tuple, tuple] = {}

        # Stats logging
        self.decision_log: List[Dict[str, Any]] = []

    def choose_lane(
        self,
        H: float,  # Roughness/entropy (0-1 normalized)
        C: float,  # Coherence/structure (0-1 normalized)
        D: float,  # Depth/scale (optional)
        budget: Budget,
    ) -> tuple:
        """
        ChatGPT's lane selection logic.

        Returns: (lane, bases, patch_k, overlap)
        """
        # Normalize H to 0-1 (entropy is 0-8 for 8-bit data)
        H_norm = min(1.0, H / 8.0)

        # Normalize C (autocorr is -1 to 1, we want 0-1)
        C_norm = (C + 1) / 2

        if budget.mode == "inference":
            # DECODE SPEED + BEHAVIORAL FIDELITY
            # Strategy: BASE_EXACT_RESIDUAL always (2 decode ops: DCT inverse + add)
            # - No sparse patches (avoid random memory access at decode)
            # - Contiguous residual storage (vectorized add)
            # - Fewer bases = faster DCT inverse
            # - Lossless reconstruction guaranteed

            # Tune bases based on structure (more structure = fewer bases needed)
            if C_norm > 0.6:
                # Highly structured: fewer bases, residual handles the rest
                return (Lane.BASE_EXACT_RESIDUAL, 128, 0, 0)
            elif C_norm > 0.3:
                # Moderate structure: balanced
                return (Lane.BASE_EXACT_RESIDUAL, 256, 0, 0)
            else:
                # Low structure / noisy: more bases to reduce residual magnitude
                return (Lane.BASE_EXACT_RESIDUAL, 384, 0, 0)

        elif budget.mode == "archive":
            # Archive: accept lower cosine for compression
            if C_norm > 0.55:
                return (Lane.FAST_GPU_QR, 256, 64, 0)
            return (Lane.FAST_GPU_QR, 128, 64, 0)

        elif budget.mode == "hybrid":
            # Hybrid: base + exact residual for perfect recon
            if C_norm > 0.55:
                return (Lane.BASE_EXACT_RESIDUAL, 512, 0, 0)
            return (Lane.BASE_EXACT_RESIDUAL, 256, 0, 0)

        # Fallback
        return (Lane.ESCAPE_RAW, 0, 0, 0)

    def decide(
        self,
        chunk: np.ndarray,
        budget: Optional[Budget] = None,
        quick_quality: Optional[float] = None,
    ) -> CodecDecision:
        """
        Decide codec parameters for a chunk.

        Args:
            chunk: The data chunk to encode
            budget: Constraints (mode, min_cos, etc.)
            quick_quality: Optional pre-computed quality score

        Returns:
            CodecDecision with all parameters
        """
        budget = budget or self.default_budget
        stats = ChunkStats.compute(chunk)

        # Use ChatGPT's lane selection
        lane, bases, patch_k, overlap = self.choose_lane(
            H=stats.entropy,
            C=stats.autocorr_lag1,
            D=stats.kurtosis,  # Using kurtosis as depth proxy
            budget=budget,
        )

        # Map lane to basis family
        if lane in (Lane.DCT_BASE, Lane.DCT_BASE_PATCH, Lane.FAST_GPU_QR):
            basis = BasisFamily.DCT
        elif lane == Lane.HYBRID_BASE_PATCH:
            basis = BasisFamily.HYBRID
        elif lane == Lane.BASE_EXACT_RESIDUAL:
            basis = BasisFamily.DCT
        else:
            basis = BasisFamily.RAW

        # Build decision
        decision = CodecDecision(
            lane=lane,
            backend=Backend.GPU_BATCHED if lane != Lane.CPU_LSTSQ else Backend.CPU,
            basis=basis,
            num_bases=bases,
            chunk_size=len(chunk),
            overlap=overlap,
            patch_k=patch_k,
            store_exact_residual=(lane == Lane.BASE_EXACT_RESIDUAL),
            store_raw=(lane == Lane.ESCAPE_RAW),
            target_cosine=budget.min_cos,
            stats=stats,
            budget=budget,
            reason=f"mode={budget.mode}, lane={lane.value}",
        )

        # Adjust based on quality feedback if available
        if quick_quality is not None and budget.mode != "hybrid":
            if quick_quality < budget.min_cos and not decision.store_raw:
                # Quality too low, escalate
                if decision.patch_k == 0:
                    decision.patch_k = 128
                    decision.reason += f"; escalated (cos={quick_quality:.3f} < {budget.min_cos})"
                elif decision.patch_k < 512:
                    decision.patch_k = min(512, decision.patch_k * 2)
                    decision.reason += f"; doubled patch"
                else:
                    decision.store_raw = True
                    decision.lane = Lane.ESCAPE_RAW
                    decision.reason += f"; escaped (patch insufficient)"

        # Log
        self.decision_log.append({
            "chunk_size": len(chunk),
            "mode": budget.mode,
            "lane": lane.value,
            "stats": {
                "entropy": stats.entropy,
                "autocorr": stats.autocorr_lag1,
                "spectral_flatness": stats.spectral_flatness,
            },
            "decision": {
                "bases": decision.num_bases,
                "patch_k": decision.patch_k,
                "store_raw": decision.store_raw,
                "store_exact": decision.store_exact_residual,
            },
            "quick_quality": quick_quality,
        })

        return decision

    def get_decision_stats(self) -> Dict[str, Any]:
        """Summary of routing decisions made."""
        if not self.decision_log:
            return {"total": 0}

        total = len(self.decision_log)
        basis_counts = {}
        escape_count = 0
        patch_count = 0

        for d in self.decision_log:
            basis = d["decision"]["basis"]
            basis_counts[basis] = basis_counts.get(basis, 0) + 1
            if d["decision"]["store_raw"]:
                escape_count += 1
            if d["decision"]["patch_k"] > 0:
                patch_count += 1

        return {
            "total": total,
            "by_basis": basis_counts,
            "escape_rate": escape_count / total,
            "patch_rate": patch_count / total,
        }


# Convenience functions

def compute_chunk_quality(chunk: np.ndarray, basis: np.ndarray, coeffs: np.ndarray) -> float:
    """Quick cosine similarity computation."""
    recon = basis @ coeffs
    return float(np.dot(chunk, recon) / (np.linalg.norm(chunk) * np.linalg.norm(recon) + 1e-10))


def extract_residual_patch(residual: np.ndarray, k: int) -> tuple:
    """
    SatGuidance-style: extract top-k residual components.

    Returns:
        (indices, values) for sparse patch storage
    """
    abs_residual = np.abs(residual)
    top_k_idx = np.argpartition(abs_residual, -k)[-k:]
    top_k_idx = top_k_idx[np.argsort(-abs_residual[top_k_idx])]  # Sort by magnitude

    return top_k_idx.astype(np.uint16), residual[top_k_idx].astype(np.float16)


def apply_residual_patch(recon: np.ndarray, indices: np.ndarray, values: np.ndarray) -> np.ndarray:
    """Apply sparse patch to reconstruction."""
    patched = recon.copy()
    patched[indices] = patched[indices] + values.astype(recon.dtype)
    return patched


# ============================================================================
# CLI Test
# ============================================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Codec Router Test")
    print("=" * 60)
    print()

    # Test on different chunk types
    np.random.seed(42)

    # 1. Random noise (like quantized weights)
    noise_chunk = np.random.randint(-128, 127, size=4096).astype(np.float32)

    # 2. Structured signal (smooth)
    t = np.linspace(0, 10*np.pi, 4096)
    structured_chunk = (np.sin(t) + 0.5*np.sin(3*t) + 0.1*np.random.randn(4096)).astype(np.float32)

    # 3. Mixed
    mixed_chunk = np.concatenate([
        np.sin(np.linspace(0, 5*np.pi, 2048)),
        np.random.randn(2048)
    ]).astype(np.float32)

    router = CodecRouter()

    for name, chunk in [("noise", noise_chunk), ("structured", structured_chunk), ("mixed", mixed_chunk)]:
        stats = ChunkStats.compute(chunk)
        decision = router.decide(chunk)

        print(f"{name.upper()} chunk:")
        print(f"  Stats: entropy={stats.entropy:.2f}, autocorr={stats.autocorr_lag1:.3f}, "
              f"spectral_flat={stats.spectral_flatness:.3f}")
        print(f"  Decision: basis={decision.basis.value}, num_bases={decision.num_bases}")
        print(f"  Reason: {decision.reason}")
        print()

    # Show routing summary
    print("Routing summary:")
    print(router.get_decision_stats())
