#!/usr/bin/env python3
"""
codebook_manager.py - Dynamic Codebook Morphing for FibPi3D Wave Regen
======================================================================

WO#58: Enables morphable weight regeneration by swapping/blending codebooks
during FibPi3D wave scheduling.

Key Concepts:
- Same compressed indices, different codebooks = different behavior
- Zone-based codebook selection (creative/sober/fast regions)
- Codebook blending for smooth transitions
- No retraining required - just swap the lookup table

Usage:
    from helix_cdc.scheduler.codebook_manager import CodebookManager

    mgr = CodebookManager()
    mgr.load("baseline", "path/to/baseline_centroids.npy")
    mgr.load("creative", "path/to/creative_centroids.npy")

    # Get codebook for a specific FibPi3D zone
    cb = mgr.get_for_zone(chunk_offset, total_size, zone_map)

    # Or blend between codebooks
    cb = mgr.blend("baseline", "creative", alpha=0.3)
"""

import numpy as np
from pathlib import Path
from typing import Dict, Optional, Tuple, Callable
import hashlib


class CodebookManager:
    """
    Manages multiple codebooks for morphable weight regeneration.

    The key insight: compressed weights are just indices into a codebook.
    Change the codebook, change the behavior - without retraining.
    """

    def __init__(self):
        self._codebooks: Dict[str, np.ndarray] = {}
        self._metadata: Dict[str, dict] = {}
        self._default: Optional[str] = None
        self._zone_map: Optional[Callable[[int, int], str]] = None

    def load(self, name: str, path: Path, set_default: bool = False) -> 'CodebookManager':
        """
        Load a codebook from .npy file.

        Args:
            name: Identifier for this codebook (e.g., "baseline", "creative")
            path: Path to .npy file containing float32 centroids
            set_default: If True, set this as the default codebook

        Returns:
            self (for chaining)
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"Codebook not found: {path}")

        cb = np.load(path).astype(np.float32)

        # Ensure sorted for consistent behavior
        cb.sort()

        self._codebooks[name] = cb
        self._metadata[name] = {
            "path": str(path),
            "bins": len(cb),
            "range": (float(cb.min()), float(cb.max())),
            "hash": hashlib.sha256(cb.tobytes()).hexdigest()[:16],
        }

        if set_default or self._default is None:
            self._default = name

        return self

    def create_uniform(self, name: str, n_bins: int = 256,
                       vmin: float = 0.0, vmax: float = 1.0,
                       set_default: bool = False) -> 'CodebookManager':
        """Create a uniform codebook (evenly spaced bins)."""
        cb = np.linspace(vmin, vmax, n_bins).astype(np.float32)
        self._codebooks[name] = cb
        self._metadata[name] = {
            "path": None,
            "bins": n_bins,
            "range": (vmin, vmax),
            "type": "uniform",
            "hash": hashlib.sha256(cb.tobytes()).hexdigest()[:16],
        }

        if set_default or self._default is None:
            self._default = name

        return self

    def create_exponential(self, name: str, n_bins: int = 16,
                           base: float = 10.0, vmin: float = 1e-4, vmax: float = 1.0,
                           set_default: bool = False) -> 'CodebookManager':
        """
        Create exponentially-spaced codebook (like the current centroid atlas).
        Good for weight magnitudes that span multiple orders of magnitude.
        """
        # Log-space the bins
        log_min = np.log10(max(vmin, 1e-10))
        log_max = np.log10(vmax)
        cb = np.logspace(log_min, log_max, n_bins).astype(np.float32)

        self._codebooks[name] = cb
        self._metadata[name] = {
            "path": None,
            "bins": n_bins,
            "range": (float(cb.min()), float(cb.max())),
            "type": "exponential",
            "hash": hashlib.sha256(cb.tobytes()).hexdigest()[:16],
        }

        if set_default or self._default is None:
            self._default = name

        return self

    def get(self, name: Optional[str] = None) -> np.ndarray:
        """Get a codebook by name (or default if None)."""
        name = name or self._default
        if name not in self._codebooks:
            raise KeyError(f"Codebook '{name}' not loaded. Available: {list(self._codebooks.keys())}")
        return self._codebooks[name]

    def blend(self, name_a: str, name_b: str, alpha: float) -> np.ndarray:
        """
        Blend two codebooks: (1-alpha)*A + alpha*B

        This enables smooth morphing between behaviors.
        Alpha=0 → pure A, Alpha=1 → pure B, Alpha=0.5 → midpoint
        """
        if name_a not in self._codebooks:
            raise KeyError(f"Codebook '{name_a}' not loaded")
        if name_b not in self._codebooks:
            raise KeyError(f"Codebook '{name_b}' not loaded")

        a = self._codebooks[name_a]
        b = self._codebooks[name_b]

        if len(a) != len(b):
            raise ValueError(f"Cannot blend codebooks of different sizes: {len(a)} vs {len(b)}")

        alpha = np.clip(alpha, 0.0, 1.0)
        blended = (1.0 - alpha) * a + alpha * b

        # Keep sorted
        blended.sort()

        return blended.astype(np.float32)

    def set_zone_map(self, zone_map: Callable[[int, int], str]) -> 'CodebookManager':
        """
        Set a zone mapping function for FibPi3D-aware codebook selection.

        The function receives (chunk_offset, total_size) and returns a codebook name.

        Example:
            def my_zone_map(offset, total):
                progress = offset / total
                if progress < 0.3:
                    return "fast"      # Early layers: quick decode
                elif progress < 0.7:
                    return "baseline"  # Middle: standard
                else:
                    return "creative"  # Late layers: more creative

            mgr.set_zone_map(my_zone_map)
        """
        self._zone_map = zone_map
        return self

    def get_for_zone(self, chunk_offset: int, total_size: int) -> np.ndarray:
        """
        Get the appropriate codebook for a FibPi3D chunk position.
        Uses the zone map if set, otherwise returns default.
        """
        if self._zone_map is None:
            return self.get(self._default)

        name = self._zone_map(chunk_offset, total_size)
        return self.get(name)

    def get_zone_name(self, chunk_offset: int, total_size: int) -> str:
        """Get the codebook name for a zone (for logging)."""
        if self._zone_map is None:
            return self._default or "default"
        return self._zone_map(chunk_offset, total_size)

    def list_codebooks(self) -> Dict[str, dict]:
        """List all loaded codebooks with metadata."""
        return {
            name: {
                **meta,
                "is_default": name == self._default
            }
            for name, meta in self._metadata.items()
        }

    def __repr__(self) -> str:
        names = list(self._codebooks.keys())
        return f"CodebookManager({len(names)} codebooks: {names}, default='{self._default}')"


# === Preset Zone Maps ===

def zone_map_thirds(offset: int, total: int) -> str:
    """Split model into thirds: early/middle/late."""
    progress = offset / max(total, 1)
    if progress < 0.33:
        return "early"
    elif progress < 0.67:
        return "middle"
    else:
        return "late"


def zone_map_attention_vs_ffn(offset: int, total: int) -> str:
    """
    Alternate between attention and FFN codebooks.
    Assumes layer structure: [attention, ffn, attention, ffn, ...]
    """
    # Crude approximation: even chunks = attention, odd = FFN
    chunk_idx = offset // 65520  # Default chunk size
    if chunk_idx % 2 == 0:
        return "attention"
    else:
        return "ffn"


def zone_map_creative_tail(offset: int, total: int, creative_threshold: float = 0.85) -> str:
    """Apply creative codebook only to the last 15% (output layers)."""
    progress = offset / max(total, 1)
    if progress >= creative_threshold:
        return "creative"
    return "baseline"


# === Factory Functions ===

def create_morphing_set(base_path: Path) -> CodebookManager:
    """
    Create a standard morphing set from a directory of codebook .npy files.

    Expected files:
        base_path/baseline.npy
        base_path/creative.npy
        base_path/fast.npy
        base_path/safe.npy (optional)
    """
    mgr = CodebookManager()

    for name in ["baseline", "creative", "fast", "safe"]:
        path = base_path / f"{name}.npy"
        if path.exists():
            mgr.load(name, path, set_default=(name == "baseline"))

    if not mgr._codebooks:
        raise FileNotFoundError(f"No codebook files found in {base_path}")

    return mgr


def create_from_existing_atlas(atlas_path: Path) -> CodebookManager:
    """
    Create a CodebookManager from an existing centroid atlas.
    Also creates perturbed variants for morphing experiments.
    """
    mgr = CodebookManager()
    mgr.load("baseline", atlas_path, set_default=True)

    baseline = mgr.get("baseline")
    n_bins = len(baseline)

    # Create perturbed variants
    # "Creative" - slightly wider distribution (more extreme values)
    creative = baseline.copy()
    center = (baseline.max() + baseline.min()) / 2
    creative = center + 1.1 * (creative - center)
    creative.sort()
    mgr._codebooks["creative"] = creative.astype(np.float32)
    mgr._metadata["creative"] = {
        "type": "perturbed",
        "base": "baseline",
        "perturbation": "expand_1.1x",
        "bins": n_bins,
    }

    # "Sober" - slightly narrower distribution (less extreme)
    sober = baseline.copy()
    sober = center + 0.9 * (sober - center)
    sober.sort()
    mgr._codebooks["sober"] = sober.astype(np.float32)
    mgr._metadata["sober"] = {
        "type": "perturbed",
        "base": "baseline",
        "perturbation": "contract_0.9x",
        "bins": n_bins,
    }

    # "Quantized" - rounded to fewer effective bins (faster but lossy)
    quantized = np.round(baseline * 100) / 100  # Round to 2 decimals
    mgr._codebooks["quantized"] = quantized.astype(np.float32)
    mgr._metadata["quantized"] = {
        "type": "perturbed",
        "base": "baseline",
        "perturbation": "round_2dec",
        "bins": n_bins,
    }

    return mgr


# === Test ===

if __name__ == "__main__":
    print("=" * 60)
    print("WO#58: Codebook Manager Test")
    print("=" * 60)

    # Create manager with synthetic codebooks
    mgr = CodebookManager()
    mgr.create_uniform("uniform", n_bins=256, vmin=0.0, vmax=1.0, set_default=True)
    mgr.create_exponential("exponential", n_bins=256, vmin=1e-4, vmax=1.0)  # Same size for blending
    mgr.create_exponential("expo16", n_bins=16, vmin=1e-4, vmax=1.0)  # 16-bin variant

    print(f"\n{mgr}")
    print("\nCodebooks:")
    for name, meta in mgr.list_codebooks().items():
        print(f"  {name}: {meta}")

    # Test blending (same-size codebooks)
    print("\nBlending test (uniform ↔ exponential, same 256 bins):")
    blend_05 = mgr.blend("uniform", "exponential", alpha=0.5)
    print(f"  Blended (alpha=0.5): {len(blend_05)} bins, range [{blend_05.min():.4f}, {blend_05.max():.4f}]")

    # Test zone mapping
    print("\nZone mapping test:")
    # Add zone codebooks
    mgr.create_uniform("early", n_bins=256, vmin=0.0, vmax=0.8)
    mgr.create_uniform("middle", n_bins=256, vmin=0.1, vmax=0.9)
    mgr.create_uniform("late", n_bins=256, vmin=0.2, vmax=1.0)
    mgr.set_zone_map(zone_map_thirds)

    for offset, total in [(0, 1000), (400, 1000), (800, 1000)]:
        zone = mgr.get_zone_name(offset, total)
        print(f"  offset={offset}, total={total} → zone='{zone}'")

    print("\n✅ CodebookManager ready for FibPi3D morphing!")
