# echo_cog/se_controller_adapter.py
"""
Adapter that wraps the historical SeController for use inside Echo Cog.

The actual SeController (from lattice_test/se_controller.py) expects:
- field: List[float] - current field values (wave amplitudes)
- prev_field: List[float] - previous field values
- window: List[List[float]] - sliding window of recent fields

Echo Cog has:
- SeState with task stats, recent Q/A history
- MMLUSample with question/choices
- WaveInfo with wave_id, phase, port

This adapter bridges the two by:
1. Converting Echo Cog state into a synthetic "field" representation
2. Mapping SeController decisions to runtime behavior hints
"""

from __future__ import annotations

import math
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

# Try to import the real SeController
_SE_CONTROLLER_AVAILABLE = False
_SeController = None
_ComputeBand = None
_RetrievalMode = None

try:
    # Add lattice_test to path if needed
    lattice_path = Path("/home/voidstr3m33/lattice_test")
    if lattice_path.exists() and str(lattice_path) not in sys.path:
        sys.path.insert(0, str(lattice_path))

    from se_controller import SeController as _SeController
    from se_controller import ComputeBand as _ComputeBand
    from se_controller import RetrievalMode as _RetrievalMode
    _SE_CONTROLLER_AVAILABLE = True
except ImportError:
    pass


@dataclass
class CogSeDecision:
    """
    What the Se controller wants the Echo Cog runtime to do for this question.

    This is a simplified view of the full SeDecision, mapped to Echo Cog concepts.
    """
    se: float = 0.0                     # Computed symbolic entropy
    h: float = 0.0                      # Shannon entropy component
    c: float = 0.0                      # Coherence component
    d: float = 0.0                      # Drift component
    band: str = "cpu"                   # "cpu", "gpu", "qpu"
    retrieval_mode: str = "off"         # "off", "readout_augment", "shell_only"
    store_flag: bool = False            # Whether to store this Q/A aggressively
    history_k_hint: Optional[int] = None  # Override history length based on Se
    max_tokens_hint: Optional[int] = None # Override max tokens based on band
    reason: str = ""                    # Why this decision was made

    def to_dict(self) -> Dict[str, Any]:
        return {
            "se": self.se,
            "h": self.h,
            "c": self.c,
            "d": self.d,
            "band": self.band,
            "retrieval_mode": self.retrieval_mode,
            "store_flag": self.store_flag,
            "history_k_hint": self.history_k_hint,
            "max_tokens_hint": self.max_tokens_hint,
            "reason": self.reason,
        }


class SeControllerAdapter:
    """
    Adapter that wraps the historical SeController for Echo Cog use.

    Converts Echo Cog state into field representations and maps
    SeController decisions back to runtime behavior hints.
    """

    def __init__(self, topology: str = "fibpi", window_size: int = 8):
        """
        Initialize the adapter.

        Args:
            topology: "fibpi", "shell", or "1d" - affects retrieval mode selection
            window_size: How many recent states to keep for Se computation
        """
        self.topology = topology
        self.window_size = window_size

        # Initialize the real SeController if available
        if _SE_CONTROLLER_AVAILABLE and _SeController is not None:
            self._controller = _SeController(window_size=window_size)
        else:
            self._controller = None

        # Sliding window of recent "fields" (synthetic from Echo Cog state)
        self._window: List[List[float]] = []
        self._prev_field: Optional[List[float]] = None

    def is_available(self) -> bool:
        """Check if the real SeController is available."""
        return self._controller is not None

    def _state_to_field(
        self,
        task_stats: Dict[str, Any],
        wave_phase: float,
        wave_port: int,
        question_entropy: float,
    ) -> List[float]:
        """
        Convert Echo Cog state into a synthetic "field" for SeController.

        The field is a vector that captures the current state:
        - Task performance (accuracy, count)
        - Wave position (phase, port)
        - Question complexity (entropy estimate)
        """
        # Build a 32-element field (matches SeController default)
        field_size = 32
        field = [0.0] * field_size

        # Encode task stats (slots 0-7)
        accuracy = task_stats.get("accuracy", 0.5)
        asked = task_stats.get("asked", 0)
        correct = task_stats.get("correct", 0)

        field[0] = accuracy
        field[1] = min(asked / 100.0, 1.0)  # Normalize count
        field[2] = min(correct / 50.0, 1.0)
        field[3] = 1.0 - accuracy  # Error rate

        # Encode wave position (slots 8-15)
        field[8] = wave_phase  # Already 0-1 normalized
        field[9] = wave_port / 8.0  # Normalize port
        field[10] = math.sin(wave_phase * 2 * math.pi)
        field[11] = math.cos(wave_phase * 2 * math.pi)

        # Encode question complexity (slots 16-23)
        field[16] = question_entropy
        field[17] = question_entropy ** 2
        field[18] = math.sqrt(max(0, question_entropy))

        # Add some structured variation (slots 24-31)
        for i in range(24, 32):
            # Fibonacci-like mixing
            if i > 25:
                field[i] = (field[i-1] + field[i-2]) / 2.0
            else:
                field[i] = (accuracy + wave_phase) / 2.0

        return field

    def _estimate_question_entropy(self, question: str, choices: List[str]) -> float:
        """
        Estimate the "entropy" of a question based on text characteristics.

        Simple heuristic:
        - Longer questions = higher entropy
        - More unique chars = higher entropy
        - Technical terms = higher entropy
        """
        full_text = question + " ".join(choices)

        # Length contribution (normalized)
        length_score = min(len(full_text) / 500.0, 1.0)

        # Unique character ratio
        unique_ratio = len(set(full_text.lower())) / max(len(full_text), 1)

        # Technical indicator (simple heuristic)
        tech_terms = ["equation", "formula", "calculate", "function", "variable",
                      "theorem", "hypothesis", "coefficient", "derivative", "integral"]
        tech_score = sum(1 for t in tech_terms if t in full_text.lower()) / len(tech_terms)

        # Combine into 0-1 entropy estimate
        entropy = 0.4 * length_score + 0.3 * unique_ratio + 0.3 * tech_score
        return min(1.0, entropy)

    def decide(
        self,
        task: str,
        task_stats: Dict[str, Any],
        wave_phase: float,
        wave_port: int,
        question: str,
        choices: List[str],
    ) -> CogSeDecision:
        """
        Make a decision for how Echo Cog should handle this question.

        Args:
            task: Task name (e.g., "anatomy")
            task_stats: Dict with "accuracy", "asked", "correct" keys
            wave_phase: Current FibPi3D wave phase (0-1)
            wave_port: Current FibPi3D port number
            question: The question text
            choices: List of answer choices

        Returns:
            CogSeDecision with routing hints for the runtime
        """
        # Estimate question entropy
        q_entropy = self._estimate_question_entropy(question, choices)

        # Convert to field representation
        current_field = self._state_to_field(task_stats, wave_phase, wave_port, q_entropy)

        # Initialize prev_field if needed
        if self._prev_field is None:
            self._prev_field = [0.0] * len(current_field)

        # Update window
        self._window.append(current_field)
        if len(self._window) > self.window_size:
            self._window.pop(0)

        # If real controller available, use it
        if self._controller is not None:
            try:
                raw_decision = self._controller.decide(
                    field=current_field,
                    prev_field=self._prev_field,
                    window=self._window,
                    topology=self.topology,
                )

                # Map to CogSeDecision
                decision = CogSeDecision(
                    se=raw_decision.se,
                    h=raw_decision.h,
                    c=raw_decision.c,
                    d=raw_decision.d,
                    band=raw_decision.band.value,
                    retrieval_mode=raw_decision.retrieval_mode.value,
                    store_flag=raw_decision.store_flag,
                    reason=raw_decision.reason,
                )

                # Map Se level to history_k and max_tokens hints
                decision.history_k_hint = self._se_to_history_k(decision.se)
                decision.max_tokens_hint = self._band_to_max_tokens(decision.band)

            except Exception as e:
                # Fallback on error
                decision = self._fallback_decision(q_entropy, f"controller_error: {e}")
        else:
            # Fallback when controller not available
            decision = self._fallback_decision(q_entropy, "controller_not_available")

        # Update prev_field for next call
        self._prev_field = current_field

        return decision

    def _se_to_history_k(self, se: float) -> int:
        """
        Map Se value to recommended history length.

        Low Se (stable) → less history needed
        High Se (volatile) → more history for context
        """
        if se < 0.05:
            return 2
        elif se < 0.10:
            return 5
        elif se < 0.15:
            return 8
        else:
            return 12

    def _band_to_max_tokens(self, band: str) -> int:
        """
        Map compute band to recommended max tokens.

        CPU (simple) → fewer tokens
        GPU (medium) → moderate tokens
        QPU (complex) → more tokens for reasoning
        """
        if band == "cpu":
            return 16
        elif band == "gpu":
            return 64
        else:  # qpu
            return 256

    def _fallback_decision(self, q_entropy: float, reason: str) -> CogSeDecision:
        """
        Generate a fallback decision when the real controller isn't available.

        Uses simple heuristics based on question entropy.
        """
        # Simple mapping: entropy → band
        if q_entropy < 0.3:
            band = "cpu"
            history_k = 3
            max_tokens = 16
        elif q_entropy < 0.6:
            band = "gpu"
            history_k = 5
            max_tokens = 64
        else:
            band = "qpu"
            history_k = 8
            max_tokens = 128

        return CogSeDecision(
            se=q_entropy,  # Use entropy as proxy for Se
            h=q_entropy,
            c=0.5,
            d=0.1,
            band=band,
            retrieval_mode="readout_augment" if q_entropy > 0.5 else "off",
            store_flag=q_entropy > 0.7,
            history_k_hint=history_k,
            max_tokens_hint=max_tokens,
            reason=f"fallback({reason}): entropy={q_entropy:.3f}",
        )

    def reset(self) -> None:
        """Reset the adapter state (clear window and prev_field)."""
        self._window.clear()
        self._prev_field = None
        if self._controller is not None:
            self._controller.reset()


# Module-level singleton for convenience
_ADAPTER_INSTANCE: Optional[SeControllerAdapter] = None


def get_se_adapter(topology: str = "fibpi") -> SeControllerAdapter:
    """Get or create the singleton SeControllerAdapter instance."""
    global _ADAPTER_INSTANCE
    if _ADAPTER_INSTANCE is None:
        _ADAPTER_INSTANCE = SeControllerAdapter(topology=topology)
    return _ADAPTER_INSTANCE


def is_se_controller_available() -> bool:
    """Check if the real SeController is available."""
    return _SE_CONTROLLER_AVAILABLE
