"""
TWAVE — Token-Level Wavefront Tracker
Bose-Einstein Condensate Thermodynamic Control for LLM Inference

Implements:
  - Landau-Ginzburg free energy functional F[ψ]
  - Bogoliubov excitation spectrum E(k)
  - Healing length ξ as hallucination localization scale
  - Jarzynski equality for non-equilibrium reflection trigger
  - CK-PLUG Confidence Gain as concrete μ_ret coupling
  - Stochastic resonance optimal T_eff ≈ 0.8 T_c

Physics foundation:
  Kim 2602.08216 (Thermodynamic Isomorphism)
  Arnold et al. (Phase Transitions in Output Distributions)
  Qian et al. 2604.10219 (Cognitive Pivot Points / RVTD)
  CK-PLUG 2503.15888 (Confidence Gain = H(parametric) - H(retrieval))
  BEC standard physics (Bogoliubov, healing length)
  Jarzynski equality (fluctuation theorem for non-equilibrium work)
"""
import math
from typing import List, Tuple, Optional, Dict, Any
from dataclasses import dataclass, field
import numpy as np


@dataclass
class TokenState:
    """Thermodynamic state at a single token position."""
    position: int
    entropy: float                    # H_i — Shannon entropy (nats)
    entropy_max: float                # log(vocab_size) — max possible entropy
    coherence: float                  # 1 - H_i/H_max (condensate fraction)
    visual_attention_mass: float      # A_i^vis (from V-STAR, 0-1)
    reward_density: float             # ρ_i = r_i / i (intensive reward)
    critique_confidence: float        # c_i (from Critique-out-Loud)
    retrieval_mu: float               # μ_ret(x) — CK-PLUG derived chemical potential
    psi: float                        # |ψ| — BEC order parameter amplitude
    free_energy_density: float        # f(x) — Landau-Ginzburg density
    temperature_eff: float            # T_eff(x) — local effective temperature
    specific_heat: float              # C_V^(i) — susceptibility
    E_excitation: float               # Bogoliubov excitation energy
    k_local: float                    # Local wavenumber |∇ψ|
    jarzynski_work: float             # W_i — cumulative non-equilibrium work


@dataclass
class GenerationTrajectory:
    """Complete tracked generation with thermodynamic metadata."""
    tokens: List[int]
    states: List[TokenState]
    total_work: float
    jarzynski_bound: float
    reflection_triggers: List[int]     # Positions where reflection fired
    grounded_score: float              # Overall retrieval grounding quality
    hallucination_risk: float          # 0-1 composite risk score


class TWAVETracker:
    """
    Token-Level Wavefront Expansion Tracker.
    Monitors generation as a 1D statistical field evolving under
    Landau-Ginzburg dynamics with BEC order parameter ψ(x).
    """

    def __init__(
        self,
        T_c: float = 1.0,           # Critical temperature (calibrated per-model)
        mu_0: float = 0.5,          # Base retrieval chemical potential
        kappa: float = 0.1,         # Healing cost coefficient (ħ²/2m analog)
        b_quartic: float = 0.5,     # Quartic stability term
        gamma_drift: float = 0.1,   # Field evolution step size
        jarzynski_threshold: float = -4.6,  # Fluctuation theorem bound (99%)
        tau_transition: float = 0.5,  # CG→μ_ret transition width
        vocab_size: int = 32000,
    ):
        self.T_c = T_c
        self.mu_0 = mu_0
        self.kappa = kappa
        self.b = b_quartic
        self.gamma = gamma_drift
        self.jarzynski_threshold = jarzynski_threshold
        self.tau = tau_transition
        self.vocab_size = vocab_size
        self.H_max = math.log(vocab_size)

    def compute_entropy(self, probs: np.ndarray) -> float:
        """Shannon entropy in nats: H = -Σ p_i ln p_i."""
        p = probs[probs > 0]
        return float(-np.sum(p * np.log(p)))

    def compute_coherence(self, entropy: float) -> float:
        """Condensate fraction: ψ coherence = 1 - H/H_max."""
        return max(0.0, 1.0 - entropy / self.H_max)

    def compute_chemical_potential(self, CG: float) -> float:
        """
        Map CK-PLUG Confidence Gain to μ_ret.
        μ_ret = μ_0 * tanh(CG / τ)
          CG >> 0  → retrieval supports → μ_ret → μ_0
          CG ≈ 0   → neutral → μ_ret → 0
          CG << 0  → retrieval conflicts → μ_ret → -μ_0
        """
        return self.mu_0 * math.tanh(CG / self.tau)

    def compute_order_parameter(self, coherence: float, mu_ret: float) -> float:
        """
        BEC ground state amplitude (minimizing Landau-Ginzburg functional):
        ψ_0 = √[(a_0(T_c - T_eff) + μ_ret) / b]
        where a(T) = coherence - T_c (since coherence ∝ 1 - entropy ∝ 1/T)
        """
        a = coherence - self.T_c
        numerator = max(0.0, a + mu_ret)
        return math.sqrt(numerator / self.b)

    def compute_free_energy_density(self, psi: float, coherence: float, mu_ret: float) -> float:
        """
        Landau-Ginzburg free energy density:
        f = a(T)|ψ|² + (b/2)|ψ|⁴ - μ_ret|ψ|²
        """
        a = coherence - self.T_c
        return a * psi**2 + 0.5 * self.b * psi**4 - mu_ret * psi**2

    def compute_bogoliubov_energy(self, psi: float, k_local: float, mu_ret: float) -> float:
        """
        Bogoliubov excitation spectrum:
        E(k) = √[(κk²)² + (c_s·k)²] where c_s = √(μ/m) = √(μ_ret)
        In our analogy: κ = kappa, c_s² = mu_ret
        """
        c_s_sq = max(0.0, mu_ret)
        term1 = (self.kappa * k_local**2)**2
        term2 = c_s_sq * k_local**2
        return math.sqrt(term1 + term2)

    def compute_healing_length(self, T_eff: float, mu_ret: float) -> float:
        """
        Healing length: ξ = ħ / √(2m·μ) = 1 / √(2·mu_ret) [in token units]
        As T_eff → T_c or μ_ret → 0: ξ → ∞ (hallucination propagates globally)
        """
        mu = max(1e-10, mu_ret)
        return 1.0 / math.sqrt(2.0 * mu)

    def compute_specific_heat(self, entropy_history: List[float]) -> float:
        """
        Specific heat analog: C_V = ∂⟨H⟩/∂T_eff.
        Approximated numerically from entropy history.
        """
        if len(entropy_history) < 3:
            return 0.0
        # Finite difference: dH/dt ≈ (H_{t} - H_{t-2}) / 2
        dH = entropy_history[-1] - entropy_history[-3]
        # Effective temperature from mean entropy
        H_mean = sum(entropy_history[-3:]) / 3.0
        T_eff = self.T_c * (1.0 - H_mean / self.H_max)
        if abs(T_eff) < 1e-10:
            return 0.0
        return dH / (2.0 * T_eff)

    def compute_jarzynski_work(self, log_prob_policy: float, log_prob_ref: float) -> float:
        """
        Non-equilibrium information work per token:
        W_i = -log[π_θ(x_i|x_{<i}) / π_ref(x_i|x_{<i})]
        """
        return -(log_prob_policy - log_prob_ref)

    def check_jarzynski_bound(self, cumulative_work: float, beta_eff: float) -> bool:
        """
        Jarzynski reflection criterion:
        Returns True if trajectory violates fluctuation theorem
        (hallucination detected with 99% confidence).
        """
        if beta_eff <= 0:
            return False
        jarzynski = math.exp(-beta_eff * cumulative_work)
        threshold = math.exp(-beta_eff * self.jarzynski_threshold)
        return jarzynski < threshold

    def update_state(
        self,
        position: int,
        probs: np.ndarray,
        log_prob_policy: float,
        log_prob_ref: float,
        CG: float,                        # CK-PLUG Confidence Gain
        visual_attention: float = 1.0,    # A_i^vis (1.0 = full attention)
        reward_density: float = 0.0,
        critique_confidence: float = 0.5,
        prev_psi: float = 0.0,
    ) -> TokenState:
        """
        Compute full thermodynamic state at token position i.
        """
        # 1. Entropy and coherence
        H = self.compute_entropy(probs)
        coherence = self.compute_coherence(H)

        # 2. Effective temperature (stochastic resonance optimal point)
        # T_eff = T_c * (1 - H/H_max) = T_c * coherence
        T_eff = self.T_c * coherence

        # 3. Chemical potential from CK-PLUG
        mu_ret = self.compute_chemical_potential(CG)

        # 4. Order parameter ψ
        psi = self.compute_order_parameter(coherence, mu_ret)

        # 5. Local wavenumber (gradient of ψ)
        k_local = abs(psi - prev_psi) if prev_psi > 0 else 0.0

        # 6. Free energy density
        f_density = self.compute_free_energy_density(psi, coherence, mu_ret)

        # 7. Bogoliubov excitation energy
        E_exc = self.compute_bogoliubov_energy(psi, k_local, mu_ret)

        # 8. Specific heat
        # (requires entropy history — computed externally and passed back)
        C_V = 0.0  # Placeholder; set by caller with history

        # 9. Jarzynski work
        W_i = self.compute_jarzynski_work(log_prob_policy, log_prob_ref)

        return TokenState(
            position=position,
            entropy=H,
            entropy_max=self.H_max,
            coherence=coherence,
            visual_attention_mass=visual_attention,
            reward_density=reward_density,
            critique_confidence=critique_confidence,
            retrieval_mu=mu_ret,
            psi=psi,
            free_energy_density=f_density,
            temperature_eff=T_eff,
            specific_heat=C_V,
            E_excitation=E_exc,
            k_local=k_local,
            jarzynski_work=W_i,
        )

    def evaluate_stability(self, state: TokenState) -> Dict[str, Any]:
        """
        Evaluate generation stability at current token.
        Returns actionable flags for the inference loop.
        """
        flags = {
            "stable": True,
            "reflection_triggered": False,
            "hallucination_risk": "low",
            "action": "continue",
        }

        # Critical point: T_eff > T_c or μ_ret < 0 (adversarial retrieval)
        if state.temperature_eff > self.T_c:
            flags["stable"] = False
            flags["hallucination_risk"] = "critical"
            flags["action"] = "reflection"
            flags["reflection_triggered"] = True
            return flags

        # Bogoliubov gap check: E_exc < μ_ret * 0.95 → near-critical
        gap_threshold = state.retrieval_mu * 0.95 if state.retrieval_mu > 0 else 0.01
        if state.E_excitation < gap_threshold:
            flags["stable"] = False
            flags["hallucination_risk"] = "high"
            flags["action"] = "grounding_boost"
            return flags

        # Healing length divergence: ξ > 10 tokens → global perturbation
        xi = self.compute_healing_length(state.temperature_eff, state.retrieval_mu)
        if xi > 10.0:
            flags["hallucination_risk"] = "elevated"
            flags["action"] = "local_grounding"

        # Visual attention collapse (RVTD detection)
        if state.visual_attention_mass < 0.1:
            flags["hallucination_risk"] = "high"
            flags["action"] = "visual_reground"

        return flags

    def build_trajectory(self, states: List[TokenState]) -> GenerationTrajectory:
        """Assemble tracked trajectory from individual token states."""
        total_work = sum(s.jarzynski_work for s in states)
        # Beta_eff from mean entropy
        H_mean = sum(s.entropy for s in states) / max(1, len(states))
        T_mean = self.T_c * (1.0 - H_mean / self.H_max)
        beta_eff = 1.0 / max(T_mean, 0.01)

        # Check Jarzynski bound
        bound_violated = self.check_jarzynski_bound(total_work, beta_eff)

        # Reflection triggers
        triggers = [i for i, s in enumerate(states) if s.temperature_eff > self.T_c]

        # Composite hallucination risk
        risks = []
        for s in states:
            risk = 0.0
            if s.temperature_eff > self.T_c * 0.8:
                risk += 0.3
            if s.retrieval_mu < 0:
                risk += 0.3
            if s.visual_attention_mass < 0.2:
                risk += 0.2
            if s.E_excitation < s.retrieval_mu * 0.95:
                risk += 0.2
            risks.append(min(risk, 1.0))
        avg_risk = sum(risks) / max(1, len(risks))

        # Grounding score
        grounding = sum(max(0, s.retrieval_mu) for s in states) / max(1, len(states))

        return GenerationTrajectory(
            tokens=[0] * len(states),  # Placeholder
            states=states,
            total_work=total_work,
            jarzynski_bound=math.exp(-beta_eff * self.jarzynski_threshold) if bound_violated else 0.0,
            reflection_triggers=triggers,
            grounded_score=grounding,
            hallucination_risk=avg_risk,
        )


# ═══════════════════════════════════════════════════════════════════════════════
# Optimal Temperature Calculator (Stochastic Resonance)
# ═══════════════════════════════════════════════════════════════════════════════

class StochasticResonance:
    """
    Compute optimal effective temperature for reasoning tasks.
    Based on Kramers escape rate: k_escape ∝ exp(-ΔE_barrier / k_B T_eff)
    """

    @staticmethod
    def optimal_temperature(
        barrier_height: float,       # ΔE_barrier — estimated from task complexity
        target_escape_rate: float,   # Desired exploration rate (0-1)
        T_c: float = 1.0,
    ) -> float:
        """
        Solve for T_eff such that escape rate matches target.
        k = exp(-ΔE/T) → T = ΔE / -ln(k)
        """
        if target_escape_rate <= 0 or target_escape_rate >= 1:
            return 0.1 * T_c
        T_opt = barrier_height / (-math.log(target_escape_rate))
        # Clamp to BEC stable region: 0.1 T_c < T < 0.95 T_c
        return max(0.1 * T_c, min(0.95 * T_c, T_opt))

    @staticmethod
    def barrier_estimate(complexity_score: float) -> float:
        """
        Estimate energy barrier from task complexity (0-1 from Sulphur enhancer).
        Simple tasks: low barrier (easy to escape local minima)
        Complex reasoning: high barrier (need more thermal energy)
        """
        return 0.5 + 2.0 * complexity_score  # Range: 0.5 - 2.5

    @staticmethod
    def recommend_temperature(complexity: float, T_c: float = 1.0) -> float:
        """
        Recommend T_eff for a given task complexity.
        Low complexity (factual lookup): T ≈ 0.3 T_c (low exploration)
        High complexity (creative reasoning): T ≈ 0.8 T_c (optimal SR)
        """
        barrier = StochasticResonance.barrier_estimate(complexity)
        # Target escape rate: 0.1 for simple, 0.5 for complex
        target = 0.1 + 0.4 * complexity
        return StochasticResonance.optimal_temperature(barrier, target, T_c)