"""
Geometric Forgetting — Decay Unused Coordinates Toward Uniform

Structural Leg L1: Without forgetting, the resonance bank grows without
bounds and loses signal-to-noise. A brain that can't prune synapses
can't think.

Mechanism:
  - Unused coordinates decay toward the uniform basin (max entropy)
  - Frequently-activated coordinates resist decay (myelinated neurons)
  - Quenched gain modulates decay resistance (what resonates with
    your disorder survives)
  - Fisher-Rao slerp ensures decay stays on manifold
  - Coordinates decayed below threshold are marked inactive, not deleted

Purity: Fisher-Rao slerp only. No Euclidean operations.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np

from .geometry import BASIN_DIM, Basin, fisher_rao_distance, slerp, to_simplex

logger = logging.getLogger(__name__)

# Uniform basin on Δ⁶³ — maximum entropy = minimum information = "forgotten"
_UNIFORM: Basin = to_simplex(np.ones(BASIN_DIM, dtype=np.float64) / BASIN_DIM)

# Below this FR distance from uniform, a coordinate is effectively forgotten
FORGOTTEN_THRESHOLD: float = 0.05


@dataclass
class DecayConfig:
    """Configuration for geometric forgetting."""

    base_rate: float = 0.001
    """Per-cycle decay rate toward uniform basin."""

    activation_protection: float = 0.5
    """How much each activation count unit shields against decay."""

    gain_modulation: float = 0.3
    """Quenched gain multiplier on decay resistance."""

    min_activations_to_protect: int = 3
    """Coordinates activated fewer times than this get no protection."""


def decay_coordinate(
    basin: Basin,
    activation_count: int,
    cycles_since_last_use: int,
    quenched_gain: float = 1.0,
    config: DecayConfig | None = None,
) -> tuple[Basin, bool]:
    """Decay a single coordinate toward uniform via FR slerp.

    Args:
        basin: Current coordinate position on Δ⁶³.
        activation_count: How many times this coordinate has been activated.
        cycles_since_last_use: Cycles since last activation.
        quenched_gain: Kernel-specific disorder gain (Pillar 3).
        config: Decay configuration.

    Returns:
        (decayed_basin, is_forgotten) — the new position and whether it
        has decayed below the forgotten threshold.
    """
    if config is None:
        config = DecayConfig()

    # Protection from activation count — neurons that fire together wire together
    if activation_count >= config.min_activations_to_protect:
        resistance = activation_count * config.activation_protection * (
            1.0 + config.gain_modulation * quenched_gain
        )
    else:
        resistance = 0.0

    effective_rate = config.base_rate / (1.0 + resistance)
    decay_amount = effective_rate * cycles_since_last_use

    # Clamp decay to [0, 1] for slerp
    decay_amount = float(np.clip(decay_amount, 0.0, 1.0))

    if decay_amount < 1e-9:
        return basin, False

    decayed = slerp(basin, _UNIFORM, decay_amount)
    dist_to_uniform = fisher_rao_distance(decayed, _UNIFORM)
    is_forgotten = dist_to_uniform < FORGOTTEN_THRESHOLD

    return decayed, is_forgotten


def decay_bank(
    coordinates: dict[int, Basin],
    activation_counts: dict[int, int],
    last_used: dict[int, int],
    current_cycle: int,
    quenched_gain: float = 1.0,
    config: DecayConfig | None = None,
) -> tuple[dict[int, Basin], set[int]]:
    """Decay all coordinates in a bank. Return updated coordinates and forgotten set.

    Args:
        coordinates: token_id → basin position.
        activation_counts: token_id → total activation count.
        last_used: token_id → cycle number of last activation.
        current_cycle: Current cycle number.
        quenched_gain: Kernel-specific disorder gain.
        config: Decay configuration.

    Returns:
        (updated_coordinates, forgotten_ids) — coordinates with decay applied
        and the set of token IDs that have decayed below threshold.
    """
    if config is None:
        config = DecayConfig()

    updated: dict[int, Basin] = {}
    forgotten: set[int] = set()

    for tid, basin in coordinates.items():
        cycles_since = current_cycle - last_used.get(tid, 0)
        count = activation_counts.get(tid, 0)

        decayed, is_forgotten = decay_coordinate(
            basin=basin,
            activation_count=count,
            cycles_since_last_use=cycles_since,
            quenched_gain=quenched_gain,
            config=config,
        )

        updated[tid] = decayed
        if is_forgotten:
            forgotten.add(tid)

    if forgotten:
        logger.debug(
            "Geometric forgetting: %d/%d coordinates forgotten",
            len(forgotten),
            len(coordinates),
        )

    return updated, forgotten
