"""
Convergence Signal Monitor & Novelty Tracker
============================================

This module provides a lightweight convergence monitor designed for
large-scale evolutionary algorithms. The monitor produces *pure signals*
describing the current optimization dynamics without making policy
decisions about exploration or exploitation.

Key design properties
---------------------
- Constant-time update (O(1)) per generation
- Constant memory footprint
- Robust to missing input signals
- Decoupled from dataset size (scales to >10^9 evaluations)

The monitor outputs smoothed diagnostic signals such as:

    - discovery_rate
    - progress_rate
    - acceptance_rate
    - stall_fraction

These signals can later be interpreted by a thermostat or control
policy that adapts mutation rates, temperature, or immigrant injection.

To support datasets growing up to 10^9 evaluations efficiently and at
O(1) memory footprints, it utilizes a flat-structure dictionary cache
with a robust memory boundary.

Author
------
Designed for scalable EZGA evolutionary optimization frameworks.
"""

from __future__ import annotations
import numpy as np
from dataclasses import dataclass
from typing import Optional, Dict
from ..core.interfaces import IConvergence


def _clip01(value: float) -> float:
    """Clamp a floating-point value to the range [0, 1]."""
    if value < 0.0:
        return 0.0
    if value > 1.0:
        return 1.0
    return float(value)


def _ema(previous: float, new_value: float, alpha: float) -> float:
    """
    Exponential moving average.

    Parameters
    ----------
    previous : float
        Previous EMA value.
    new_value : float
        New observed value.
    alpha : float
        Smoothing factor in (0, 1].

    Returns
    -------
    float
        Updated EMA value.

    Notes
    -----
    EMA provides stable smoothing with constant memory:

    EMA_t = (1 - alpha) * EMA_{t-1} + alpha * x_t
    """
    if alpha <= 0.0:
        return previous
    if alpha >= 1.0:
        return new_value
    return (1 - alpha) * previous + alpha * new_value


@dataclass
class ConvergenceSignals:
    """
    Container holding the current convergence signals.

    These signals are *observables only* and do not contain policy
    decisions. They can be consumed by thermostats or adaptive operators.

    Attributes
    ----------
    generation : int
        Current generation index.

    discovery_rate : float
        Exponential moving average of newly discovered structures.

    progress_rate : float
        Exponential moving average of improvement events.

    acceptance_rate : float
        Smoothed transition acceptance ratio.

    stall_fraction : float
        Normalized stagnation measure in [0, 1].

    improvement : bool
        Raw improvement event signal.

    new_structures : int
        Number of novel structures discovered.

    accepted_structures : int
        Number of accepted individuals after transition.

    proposed_structures : int
        Number of proposed candidates.
    """
    generation: int = 0
    discovery_rate: float = 0.0
    progress_rate: float = 0.0
    acceptance_rate: float = 0.0
    stall_fraction: float = 0.0

    improvement: bool = False
    new_structures: int = 0
    accepted_structures: int = 0
    proposed_structures: int = 0


class Convergence(IConvergence):
    """
    Convergence monitor producing optimization dynamics signals and tracking features.

    This class tracks generating signals natively compatible with `engine.py`, 
    whilst supporting memory-bounded cache to scale cleanly up to 10^9 structures.

    The monitor does **not** interpret these signals. Instead, it exposes
    them so that higher-level policies (e.g. thermostats) can adapt the
    evolutionary dynamics.

    Default behavior
    ----------------
    By default, convergence is declared if no objective improvement has
    occurred for `stall_threshold` generations.

    Complexity
    ----------
    Time complexity per update:
        O(1) (When track_features is False)

    Memory complexity:
        O(1) (Bounded effectively when track_features is True)

    This makes the monitor safe for runs containing billions of
    evaluations.
    """

    def __init__(
        self,
        stall_threshold: Optional[int] = None,
        objective_threshold: float = 1e-4,
        alpha_discovery: float = 0.05,
        alpha_progress: float = 0.05,
        alpha_acceptance: float = 0.1,
        max_cache_size: int = 5_000_000,
        track_features: bool = False,
    ):
        """
        Initialize the highly efficient convergence monitor.

        Parameters
        ----------
        stall_threshold : int
            Number of generations without improvement before declaring convergence. 
            If None, convergence effectively never halts early.
        objective_threshold : float
            Strictness threshold for what is defined as an objective improvement.
        max_cache_size : int
            Boundary dict limit. Handles worst-case O(N) memory leak linearly
            without sacrificing O(1) EMA operations if datasets scale > 10^9 elements.
        track_features : bool
            If true, tracks the objective improvements per individual feature vector (slow). 
            If false, tracks the absolute global objective minimum (O(1) ultra-fast).
        """
        # We interpret `None` as infinity internally
        self._stall_threshold = float('inf') if stall_threshold is None else int(stall_threshold)
        self._objective_threshold = float(objective_threshold)
        self._track_features = bool(track_features)
        
        self._alpha_discovery = float(alpha_discovery)
        self._alpha_progress = float(alpha_progress)
        self._alpha_acceptance = float(alpha_acceptance)

        self._max_cache_size = int(max_cache_size)

        # Internal state
        self._generation = 0
        self._stall_counter = 0

        self._ema_discovery = 0.0
        self._ema_progress = 0.0
        self._ema_acceptance = 0.0

        self._converged = False
        self._improvement_found = False

        self.signals = ConvergenceSignals()
        
        # Extremely fast 1D dictionary map. Structure: Tuple -> Array 
        self._best_objectives = {}
        
        # O(1) global fallback
        self._global_best = None

    def _evict_cache_if_needed(self):
        """
        Memory bounding logic. Prevents dictionary from scaling beyond max scale
        without incurring LRU per-access cost. Drops ~25% oldest entries.
        """
        if len(self._best_objectives) > self._max_cache_size:
            keys = list(self._best_objectives.keys())
            # Keep the more recent iterations (approximate LRU by dictionary insertion order)
            cutoff = int(len(keys) * 0.25)
            
            new_dict = {}
            for k in keys[cutoff:]:
                new_dict[k] = self._best_objectives[k]
                
            self._best_objectives = new_dict

    def check(
        self,
        generation: int,
        objectives,
        features,
        accepted_structures: int = 0,
        proposed_structures: int = 0,
    ) -> Dict[str, object]:
        """
        Updates convergence signals by processing raw targets and novelty features natively.
        Runs in extremely fast O(B) dictionary time constraint.
        """
        self._generation = generation

        # Ensure objectives and features are at least 1D arrays
        objectives_arr = np.atleast_1d(objectives)
        features_arr = np.atleast_1d(features)

        if objectives_arr.ndim == 1:
            objectives_arr = objectives_arr.reshape(-1, 1)
        if features_arr.ndim == 1:
            features_arr = features_arr.reshape(-1, 1)

        n_structures = objectives_arr.shape[0]

        improvement_found = False
        new_structures = 0

        if not self._track_features:
            # Ultra-fast O(N numpy) / O(1 Python) convergence check
            current_min = np.min(objectives_arr)
            if self._global_best is None:
                self._global_best = current_min
                improvement_found = True
            elif current_min < self._global_best - self._objective_threshold:
                self._global_best = current_min
                improvement_found = True
        else:
            # High performance dictionary map
            for i in range(n_structures):
                feature_key = tuple(features_arr[i])
                current_obj = objectives_arr[i]
                
                prev_best = self._best_objectives.get(feature_key)
                if prev_best is not None:
                    if np.any(current_obj < prev_best - self._objective_threshold):
                        improvement_found = True
                        self._best_objectives[feature_key] = np.minimum(prev_best, current_obj)
                else:
                    new_structures += 1
                    improvement_found = True
                    self._best_objectives[feature_key] = current_obj.copy()

            # OOM safety checks
            self._evict_cache_if_needed()

        self._improvement_found = improvement_found

        if improvement_found:
            self._stall_counter = 0
        else:
            self._stall_counter = min(self._stall_threshold, self._stall_counter + 1)

        # Update smoothed signals
        self._ema_discovery = _ema(self._ema_discovery, float(new_structures), self._alpha_discovery)
        self._ema_progress = _ema(self._ema_progress, 1.0 if improvement_found else 0.0, self._alpha_progress)

        if proposed_structures > 0:
            acceptance = accepted_structures / proposed_structures
        else:
            acceptance = 1.0 if accepted_structures > 0 else 0.0

        self._ema_acceptance = _ema(self._ema_acceptance, acceptance, self._alpha_acceptance)

        stall_fraction = _clip01(self._stall_counter / self._stall_threshold)
        self._converged = self._stall_counter >= self._stall_threshold

        # Update signal container
        self.signals = ConvergenceSignals(
            generation=generation,
            discovery_rate=self._ema_discovery,
            progress_rate=self._ema_progress,
            acceptance_rate=self._ema_acceptance,
            stall_fraction=stall_fraction,
            improvement=improvement_found,
            new_structures=new_structures,
            accepted_structures=accepted_structures,
            proposed_structures=proposed_structures,
        )
        return {
            "converge": self._converged,
            "improvement_found": self._improvement_found,
            "stagnation": self._stall_counter,
            "signals": self.signals,
            "stall_count_objetive": self._stall_counter, # Backwards compatibility
        }

    # ---------------------------------------------------------
    # Required interface for EZGA engine
    # ---------------------------------------------------------

    def is_converge(self) -> bool:
        return self._converged
        
    def improvement_found(self) -> bool:
        return self._improvement_found

    def get_stagnation(self) -> int:
        return self._stall_counter

    # ---------------------------------------------------------
    # Convenience properties
    # ---------------------------------------------------------

    @property
    def discovery_rate(self) -> float:
        return self._ema_discovery

    @property
    def progress_rate(self) -> float:
        return self._ema_progress

    @property
    def acceptance_rate(self) -> float:
        return self._ema_acceptance

    @property
    def stall_fraction(self) -> float:
        return _clip01(self._stall_counter / self._stall_threshold)

    # ---------------------------------------------------------
    # Reset utility
    # ---------------------------------------------------------

    def reset(self) -> None:
        self._generation = 0
        self._stall_counter = 0

        self._ema_discovery = 0.0
        self._ema_progress = 0.0
        self._ema_acceptance = 0.0

        self._converged = False
        self._improvement_found = False
        
        self._best_objectives.clear()
        self._global_best = None
        self.signals = ConvergenceSignals()