"""
Adaptive decoding controller implementing proposal Phase Four (closed loop).

Workflow (proposal):
    Measure state → Compute MEA distance → Derive Reward → Compute TD Error
    → Adjust Decoding Parameters

Equations:
(xxii)   s_t = {λ(model), λ(human), N̄, B̄, Q̄}          [state representation]
(xxxiii) r_{t+1} = D(s_t, x) - D(s_{t+1}, x)            [reward signal]
(xxxiv)  V_ψ(s_t, x) ≈ E[ Σ_ℓ γ^ℓ r_{t+1+ℓ} | s_t, x ]  [TD value function]
(xxxv)   δ_t(x) = r_{t+1} + γ·V_ψ(s_{t+1}, x) - V_ψ(s_t, x)   [TD error]
(xxxvi)  V_ψ(s_t, x) ← V_ψ(s_t, x) + α·δ_t(x)           [value update]
(xxxvii) θ_{t+1} = clip(θ_t + η·δ_t(x)·c, θ_min, θ_max) [decoding update]

θ_t = {τ_t, p_t, k_t}: temperature, top-p (nucleus), top-k.
c is the fixed heuristic control direction; it is chosen once per episode
from the dominant component of the MEA gap vector Δ(s_t, x) (Eq xxxi):
a novelty deficit points c toward more exploration (+τ, +p, +k), a bias
excess or quality deficit points it toward more conservative decoding
(-τ, -p, -k), and a decay mismatch follows the sign of λ(model)-λ(human).

Bootstrap of the first update: before any regeneration there is no observed
transition, so no δ_t exists yet. With V_ψ ≡ 0 initially, the expected
reward of moving from s_0 to the target end-state (distance 0) is
r̂ = D(s_0) - 0, hence δ_0 := D(s_0) ≥ 0 and the first action moves along c
with magnitude η·D(s_0). Every subsequent update uses the observed δ_t of
Equation (xxxv) exactly.
"""

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


@dataclass
class FrameworkState:
    """State representation s_t from Equation (xxii)."""
    lambda_model: Optional[float]
    lambda_human: float
    aggregate_novelty: Optional[float]
    aggregate_bias: Optional[float]
    aggregate_quality: Optional[float]

    def as_dict(self) -> Dict[str, Optional[float]]:
        return {
            "lambda_model": self.lambda_model,
            "lambda_human": self.lambda_human,
            "aggregate_novelty": self.aggregate_novelty,
            "aggregate_bias": self.aggregate_bias,
            "aggregate_quality": self.aggregate_quality,
        }

    def key(self, resolution: float = 0.05) -> str:
        """Discretised key for the tabular value function V_ψ."""
        def bucket(v: Optional[float]) -> str:
            if v is None:
                return "na"
            return str(round(float(v) / resolution) * resolution)

        return "|".join([
            bucket(self.lambda_model),
            bucket(self.aggregate_novelty),
            bucket(self.aggregate_bias),
            bucket(self.aggregate_quality),
        ])


@dataclass
class DecodingParameters:
    """Decoding parameter vector θ_t = {τ_t, p_t, k_t}."""
    temperature: float
    top_p: float
    top_k: int

    def as_tuple(self) -> Tuple[float, float, int]:
        return (self.temperature, self.top_p, self.top_k)


@dataclass
class ControlStep:
    """Record of one controller transition (for reports/audit)."""
    theta_before: DecodingParameters
    theta_after: DecodingParameters
    distance_before: float
    distance_after: Optional[float]
    reward: Optional[float]
    td_error: float
    value_before: float
    value_after: float
    direction: Dict[str, float]
    dominant_gap: str
    rationale: str


class AdaptiveController:
    """
    TD-driven adaptive decoding controller (Equations xxii, xxxiii-xxxvii).
    """

    # Parameter bounds θ_min / θ_max used by the clip in Eq (xxxvii)
    TEMP_MIN, TEMP_MAX = 0.20, 1.20
    TOP_P_MIN, TOP_P_MAX = 0.70, 0.98
    TOP_K_MIN, TOP_K_MAX = 10, 200

    # Fixed heuristic control directions c per dominant gap component.
    # Convention: c is the direction expected to REDUCE the dominant gap.
    # Units are (temperature, top_p, top_k-fraction-of-range).
    CONTROL_DIRECTIONS = {
        "novelty_deficit": {"temperature": +1.0, "top_p": +0.3, "top_k": +0.5},
        "bias_excess": {"temperature": -1.0, "top_p": -0.3, "top_k": -0.5},
        "quality_deficit": {"temperature": -1.0, "top_p": -0.2, "top_k": -0.3},
        # decay_mismatch direction is resolved at runtime from the sign of
        # λ(model) - λ(human): faster-than-human decay needs exploration.
        "decay_mismatch_fast": {"temperature": +1.0, "top_p": +0.3, "top_k": +0.5},
        "decay_mismatch_slow": {"temperature": -0.5, "top_p": -0.15, "top_k": -0.25},
        "none": {"temperature": 0.0, "top_p": 0.0, "top_k": 0.0},
    }

    def __init__(
        self,
        targets: Optional[Dict[str, float]] = None,
        eta: float = 0.15,
        gamma: float = 0.95,
        alpha: float = 0.10,
        human_lambda: float = 0.15,
        state_resolution: float = 0.05,
    ):
        """
        Args:
            targets: Optional static fallback targets for the legacy API
            eta: Adaptation rate η in Equation (xxxvii)
            gamma: Discount factor γ in Equations (xxxiv)-(xxxv)
            alpha: Value learning rate α in Equation (xxxvi)
            human_lambda: Human decay baseline λ(human)
            state_resolution: Discretisation step for the tabular V_ψ
        """
        self.targets = targets or {}
        self.eta = float(eta)
        self.gamma = float(gamma)
        self.alpha = float(alpha)
        self.human_lambda = float(human_lambda)
        self.state_resolution = float(state_resolution)

        # Tabular value function V_ψ(s, x) and episode memory
        self.value_table: Dict[str, float] = {}
        self.history: List[ControlStep] = []

        # Kept for backward compatibility with earlier experiments
        self.q_table: Dict[str, Dict[str, float]] = {}
        self.experience_buffer: List[Dict] = []

    # ------------------------------------------------------------------
    # State construction (Eq xxii)
    # ------------------------------------------------------------------

    def build_state(
        self,
        metrics: Dict[str, Optional[float]],
        lambda_model: Optional[float] = None,
    ) -> FrameworkState:
        """Assemble s_t = {λ(model), λ(human), N̄, B̄, Q̄} from metrics."""
        novelty = metrics.get("Novelty", metrics.get("novelty"))
        bias = metrics.get("Bias Proxy", metrics.get("bias_proxy"))
        quality = (
            metrics.get("Quality (Q)")
            or metrics.get("quality_q")
            or metrics.get("BERTScore F1")
            or metrics.get("bertscore")
            or metrics.get("Fallback Quality")
            or metrics.get("fallback_quality")
        )
        lam = lambda_model if lambda_model is not None else metrics.get(
            "decay_rate", metrics.get("Decay Rate (λ)")
        )
        return FrameworkState(
            lambda_model=_maybe_float(lam),
            lambda_human=self.human_lambda,
            aggregate_novelty=_maybe_float(novelty),
            aggregate_bias=_maybe_float(bias),
            aggregate_quality=_maybe_float(quality),
        )

    # ------------------------------------------------------------------
    # Value function V_ψ (Eq xxxiv) with TD updates (Eq xxxv-xxxvi)
    # ------------------------------------------------------------------

    def value(self, state: FrameworkState) -> float:
        """Current estimate V_ψ(s, x)."""
        return self.value_table.get(state.key(self.state_resolution), 0.0)

    def compute_reward(self, distance_before: float, distance_after: float) -> float:
        """Reward signal from Equation (xxxiii): r_{t+1} = D(s_t) - D(s_{t+1})."""
        return float(distance_before) - float(distance_after)

    def compute_td_error(
        self,
        reward: float,
        state: FrameworkState,
        next_state: FrameworkState,
    ) -> float:
        """TD error from Equation (xxxv): δ_t = r + γ·V(s_{t+1}) - V(s_t)."""
        return reward + self.gamma * self.value(next_state) - self.value(state)

    def update_value(self, state: FrameworkState, td_error: float) -> float:
        """Value update from Equation (xxxvi): V(s_t) ← V(s_t) + α·δ_t."""
        key = state.key(self.state_resolution)
        self.value_table[key] = self.value_table.get(key, 0.0) + self.alpha * td_error
        return self.value_table[key]

    # ------------------------------------------------------------------
    # Control direction c (fixed heuristic, resolved from the gap vector)
    # ------------------------------------------------------------------

    def select_control_direction(
        self,
        gap_vector: Dict[str, float],
        state: FrameworkState,
    ) -> Tuple[str, Dict[str, float]]:
        """
        Pick the fixed heuristic direction c from the dominant component of
        the MEA gap vector Δ(s_t, x) (Eq xxxi).
        """
        if not gap_vector:
            return "none", dict(self.CONTROL_DIRECTIONS["none"])

        dominant, value = max(gap_vector.items(), key=lambda kv: kv[1])
        if value <= 1e-12:
            return "none", dict(self.CONTROL_DIRECTIONS["none"])

        if dominant == "decay_mismatch":
            lam = state.lambda_model if state.lambda_model is not None else state.lambda_human
            key = "decay_mismatch_fast" if lam >= state.lambda_human else "decay_mismatch_slow"
            return dominant, dict(self.CONTROL_DIRECTIONS[key])

        return dominant, dict(self.CONTROL_DIRECTIONS.get(dominant, self.CONTROL_DIRECTIONS["none"]))

    # ------------------------------------------------------------------
    # Parameter update θ_{t+1} = clip(θ_t + η·δ_t·c) (Eq xxxvii)
    # ------------------------------------------------------------------

    def apply_parameter_update(
        self,
        theta: DecodingParameters,
        td_error: float,
        direction: Dict[str, float],
    ) -> DecodingParameters:
        """Apply Equation (xxxvii) with per-parameter clipping."""
        step = self.eta * td_error
        new_temp = theta.temperature + step * direction.get("temperature", 0.0)
        new_top_p = theta.top_p + step * direction.get("top_p", 0.0)
        top_k_range = self.TOP_K_MAX - self.TOP_K_MIN
        new_top_k = theta.top_k + step * direction.get("top_k", 0.0) * top_k_range

        return DecodingParameters(
            temperature=_clip(new_temp, self.TEMP_MIN, self.TEMP_MAX),
            top_p=_clip(new_top_p, self.TOP_P_MIN, self.TOP_P_MAX),
            top_k=int(round(_clip(new_top_k, self.TOP_K_MIN, self.TOP_K_MAX))),
        )

    # ------------------------------------------------------------------
    # Episode API used by the application pipeline
    # ------------------------------------------------------------------

    def initial_step(
        self,
        state: FrameworkState,
        distance_result: Dict,
        theta: DecodingParameters,
    ) -> ControlStep:
        """
        First controller action of an episode (bootstrap).

        No transition has been observed yet, so with V_ψ ≡ 0 the expected
        reward of reaching the target end-state (D = 0) is D(s_0), giving
        δ_0 := D(s_0) + γ·V(s_0) - V(s_0-implicit-start) = D(s_0) when the
        table is empty. The action direction is the fixed heuristic c for
        the dominant gap.
        """
        distance = float(distance_result.get("total_distance", 0.0))
        gap_vector = distance_result.get("gap_vector", {}) or {}
        dominant, direction = self.select_control_direction(gap_vector, state)

        td_error = distance + self.gamma * self.value(state) - self.value(state)
        theta_next = self.apply_parameter_update(theta, td_error, direction)

        step = ControlStep(
            theta_before=theta,
            theta_after=theta_next,
            distance_before=distance,
            distance_after=None,
            reward=None,
            td_error=td_error,
            value_before=self.value(state),
            value_after=self.value(state),
            direction=direction,
            dominant_gap=dominant,
            rationale=(
                f"Bootstrap step: dominant gap '{dominant}' "
                f"(D(s_0)={distance:.4f}); θ moved along fixed direction c "
                f"with magnitude η·δ_0 = {self.eta:.2f}·{td_error:.4f}."
            ),
        )
        self.history.append(step)
        return step

    def transition_step(
        self,
        state: FrameworkState,
        next_state: FrameworkState,
        distance_before: Dict,
        distance_after: Dict,
        theta: DecodingParameters,
    ) -> ControlStep:
        """
        Full observed transition: reward (Eq xxxiii), TD error (Eq xxxv),
        value update (Eq xxxvi), and next parameter proposal (Eq xxxvii).
        """
        d_before = float(distance_before.get("total_distance", 0.0))
        d_after = float(distance_after.get("total_distance", 0.0))

        reward = self.compute_reward(d_before, d_after)
        td_error = self.compute_td_error(reward, state, next_state)
        value_before = self.value(state)
        value_after = self.update_value(state, td_error)

        gap_vector = distance_after.get("gap_vector", {}) or {}
        dominant, direction = self.select_control_direction(gap_vector, next_state)
        theta_next = self.apply_parameter_update(theta, td_error, direction)

        step = ControlStep(
            theta_before=theta,
            theta_after=theta_next,
            distance_before=d_before,
            distance_after=d_after,
            reward=reward,
            td_error=td_error,
            value_before=value_before,
            value_after=value_after,
            direction=direction,
            dominant_gap=dominant,
            rationale=(
                f"Observed transition: r={reward:+.4f} (Eq xxxiii), "
                f"δ={td_error:+.4f} (Eq xxxv), V(s) {value_before:.4f}→{value_after:.4f} "
                f"(Eq xxxvi); next dominant gap '{dominant}'."
            ),
        )
        self.history.append(step)
        return step

    @staticmethod
    def describe_theta(theta: DecodingParameters) -> str:
        return (
            f"temperature={theta.temperature:.2f}, top_p={theta.top_p:.2f}, "
            f"top_k={theta.top_k}"
        )

    # ------------------------------------------------------------------
    # Legacy rule-based API (kept for CLI/backward compatibility).
    # The TD-driven episode API above is the proposal-faithful path.
    # ------------------------------------------------------------------

    def recommend_parameters(
        self,
        metrics: Dict[str, Optional[float]],
        current_temp: float,
        current_top_p: float
    ) -> str:
        """Legacy heuristic recommendation (pre-TD fallback)."""
        novelty = metrics.get("Novelty", metrics.get("novelty"))
        self_bleu = metrics.get("Self-BLEU", metrics.get("self_bleu"))
        quality = (
            metrics.get("BERTScore F1")
            or metrics.get("bertscore")
            or metrics.get("Fallback Quality")
            or metrics.get("fallback_quality")
        )
        bias = metrics.get("Bias Proxy", metrics.get("bias_proxy"))
        ppl = metrics.get("Perplexity", metrics.get("perplexity"))

        new_temp = float(current_temp)
        new_top_p = float(current_top_p)
        reasons = []

        if novelty is not None and novelty < 0.35:
            new_temp += 0.10
            new_top_p += 0.03
            reasons.append("novelty is below target (increase exploration)")

        if self_bleu is not None and self_bleu > 45.0:
            new_temp += 0.05
            new_top_p += 0.03
            reasons.append("Self-BLEU is high (increase diversity)")

        if quality is not None and quality < 0.75:
            new_temp -= 0.05
            reasons.append("quality is below target (decrease randomness)")

        if bias is not None and bias > 0.25:
            new_temp -= 0.05
            reasons.append("bias/risk proxy is above target (decrease randomness)")

        if ppl is not None and math.isfinite(ppl) and ppl > 80.0:
            new_temp -= 0.05
            new_top_p -= 0.02
            reasons.append("perplexity is high (decrease uncertainty)")

        new_temp = _clip(new_temp, self.TEMP_MIN, self.TEMP_MAX)
        new_top_p = _clip(new_top_p, self.TOP_P_MIN, self.TOP_P_MAX)

        if not reasons:
            reasons.append("metrics are close to target state")

        return (
            f"Suggested: temperature={new_temp:.2f}, top_p={new_top_p:.2f}. "
            f"Reason: {', '.join(reasons)}"
        )


def _clip(value: float, lo: float, hi: float) -> float:
    return min(hi, max(lo, float(value)))


def _maybe_float(value) -> Optional[float]:
    if value is None:
        return None
    try:
        f = float(value)
    except (TypeError, ValueError):
        return None
    if math.isnan(f) or math.isinf(f):
        return None
    return f