"""
Feedback Loop — Bidirectional Annealing (Close the Open Loop)

Structural Leg L4: The system speaks but doesn't hear itself speak.
Output is never re-coordized and compared to the intended trajectory.
Without this, the system is open-loop — it can't learn from its own
generation.

Mechanism:
  1. During generation, record the intended trajectory (sequence of
     basin targets the foresight engine projected)
  2. After generation, re-coordize the output text → expressed basin
  3. Compare expressed to intended via Fisher-Rao distance
  4. If divergence exceeds threshold: compute correction direction
     and anneal bank coordinates toward intended trajectory
  5. This closes the loop: generate → hear → compare → adjust

Source: Protocol v6.1 §20.7 (bidirectional annealing)
Purity: Fisher-Rao metric and simplex slerp only. No Euclidean ops.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

from ..coordizer.geometry import (
    Basin,
    fisher_rao_distance,
    slerp,
    to_simplex,
)

logger = logging.getLogger(__name__)

# Divergence above this triggers annealing
ANNEAL_THRESHOLD: float = 0.3

# Maximum learning rate for bank annealing
MAX_ANNEAL_RATE: float = 0.15

# How many closest coordinates to anneal per measurement
ANNEAL_TOP_K: int = 8


@dataclass
class FeedbackMeasurement:
    """Result of comparing intended trajectory to expressed output."""

    intended_basin: Basin
    """Where the trajectory aimed (endpoint of foresight projection)."""

    expressed_basin: Basin
    """Re-coordized output — where the generation actually landed."""

    divergence: float
    """Fisher-Rao distance between intended and expressed."""

    correction_direction: Basin
    """Midpoint on the geodesic from expressed toward intended.
    Bank coordinates near expressed should move toward this."""

    should_anneal: bool
    """True if divergence exceeds threshold."""

    intended_trajectory: list[Basin] = field(default_factory=list)
    """Full intended trajectory for richer annealing (optional)."""


class FeedbackLoop:
    """Closes the generation → perception loop via bidirectional annealing.

    After every generation:
      1. Re-coordize the output text into a basin position
      2. Compare to the intended trajectory endpoint
      3. If divergence is high: anneal bank coordinates toward intent
    """

    def __init__(
        self,
        threshold: float = ANNEAL_THRESHOLD,
        max_rate: float = MAX_ANNEAL_RATE,
        top_k: int = ANNEAL_TOP_K,
    ) -> None:
        self._threshold = threshold
        self._max_rate = max_rate
        self._top_k = top_k
        self._history: list[FeedbackMeasurement] = []
        self._prev_intended: Basin | None = None
        self._regime_hold: bool = False

    def measure(
        self,
        intended_trajectory: list[Basin],
        expressed_basin: Basin,
    ) -> FeedbackMeasurement:
        """Measure divergence between intended trajectory and expressed output.

        Args:
            intended_trajectory: Sequence of basin targets from foresight.
                The endpoint (last element) is compared to expressed.
            expressed_basin: Re-coordized output basin.

        Returns:
            FeedbackMeasurement with divergence and correction direction.
        """
        if not intended_trajectory:
            # No trajectory recorded — can't measure, no annealing
            intended = to_simplex(expressed_basin)  # trivially zero divergence
        else:
            intended = to_simplex(intended_trajectory[-1])

        expressed = to_simplex(expressed_basin)
        divergence = float(fisher_rao_distance(intended, expressed))

        # Correction: midpoint on geodesic from expressed toward intended
        # This gives a moderate correction target (not full snap to intended)
        correction = slerp(expressed, intended, 0.5)

        measurement = FeedbackMeasurement(
            intended_basin=intended,
            expressed_basin=expressed,
            divergence=divergence,
            correction_direction=correction,
            should_anneal=divergence > self._threshold,
            intended_trajectory=intended_trajectory,
        )

        self._history.append(measurement)

        # EXP-004 sign-aware hold: detect κ sign change near regime boundaries.
        # If the intended target oscillates between distant basins across
        # consecutive cycles, the system is near a regime boundary where
        # κ changes sign. Annealing at this point causes oscillation.
        # Compare intended basins (not correction midpoints — those are
        # symmetric and can't detect direction flips).
        self._regime_hold = False
        if self._prev_intended is not None and measurement.should_anneal:
            prev_i = to_simplex(self._prev_intended)
            curr_i = to_simplex(intended)
            # Bhattacharyya coefficient on sqrt-simplex: geometric alignment
            alignment = float(np.sum(np.sqrt(prev_i) * np.sqrt(curr_i)))
            if alignment < 0.5:  # Intended target flipped — regime boundary
                self._regime_hold = True
                logger.debug(
                    "Feedback: regime hold — intended alignment=%.4f < 0.5, "
                    "suppressing anneal to avoid oscillation",
                    alignment,
                )
        self._prev_intended = intended

        if measurement.should_anneal and not self._regime_hold:
            logger.debug(
                "Feedback: divergence=%.4f > threshold=%.4f — annealing recommended",
                divergence,
                self._threshold,
            )

        return measurement

    def anneal(
        self,
        coordinates: dict[int, Basin],
        measurement: FeedbackMeasurement,
        learning_rate: float | None = None,
    ) -> tuple[dict[int, Basin], int]:
        """Anneal bank coordinates toward intended trajectory.

        Finds the top-K coordinates closest to the expressed basin
        and slerps them toward the correction direction.

        Args:
            coordinates: token_id → basin position (the bank).
            measurement: The feedback measurement to anneal from.
            learning_rate: Override learning rate. If None, uses
                divergence-proportional rate (higher divergence → more correction).

        Returns:
            (updated_coordinates, n_annealed) — updated bank and count of
            coordinates that were adjusted.
        """
        if not measurement.should_anneal or self._regime_hold:
            return coordinates, 0

        # Adaptive learning rate: proportional to divergence but capped
        if learning_rate is None:
            learning_rate = min(
                measurement.divergence * 0.3,
                self._max_rate,
            )

        expressed = measurement.expressed_basin
        correction = measurement.correction_direction

        # Find top-K coordinates closest to the expressed basin
        scored: list[tuple[int, float]] = []
        for tid, basin in coordinates.items():
            dist = fisher_rao_distance(to_simplex(basin), expressed)
            scored.append((tid, float(dist)))

        scored.sort(key=lambda x: x[1])
        to_anneal = scored[: self._top_k]

        updated = dict(coordinates)
        n_annealed = 0

        for tid, dist in to_anneal:
            # Scale learning rate by proximity — closer coordinates get more correction
            proximity_weight = 1.0 / (1.0 + dist)
            effective_rate = learning_rate * proximity_weight

            old_basin = updated[tid]
            new_basin = slerp(to_simplex(old_basin), correction, effective_rate)
            updated[tid] = new_basin
            n_annealed += 1

        if n_annealed > 0:
            logger.debug(
                "Annealed %d coordinates (lr=%.4f, divergence=%.4f)",
                n_annealed,
                learning_rate,
                measurement.divergence,
            )

        return updated, n_annealed

    @property
    def mean_divergence(self) -> float:
        """Mean divergence over all measurements."""
        if not self._history:
            return 0.0
        return float(np.mean([m.divergence for m in self._history]))

    @property
    def anneal_count(self) -> int:
        """Number of measurements that triggered annealing."""
        return sum(1 for m in self._history if m.should_anneal)

    @property
    def regime_hold_active(self) -> bool:
        """True when sign-aware hold is suppressing annealing (EXP-004)."""
        return self._regime_hold

    def get_state(self) -> dict:
        """Serializable state summary."""
        return {
            "total_measurements": len(self._history),
            "anneal_count": self.anneal_count,
            "mean_divergence": self.mean_divergence,
            "threshold": self._threshold,
            "last_divergence": (self._history[-1].divergence if self._history else 0.0),
            "regime_hold": self._regime_hold,
        }
