"""
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 (when track_features=False)
- Constant memory footprint
- Robust to missing input signals
- Decoupled from dataset size (scales to >10^10 evaluations)

The monitor outputs smoothed diagnostic signals such as:
    - discovery_rate: Volume of novel structures.
    - progress_rate: Frequency of objective improvement.
    - acceptance_rate: Quality of proposer/kernel transitions.

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, Any, Tuple
from ..core.interfaces import IConvergence


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

    Args:
        value (float): The input value to clamp.

    Returns:
        float: The value restricted to [0.0, 1.0].
    """
    return float(np.clip(value, 0.0, 1.0))


def _ema(previous: float, new_value: float, alpha: float) -> float:
    """
    Compute the Exponential Moving Average (EMA).

    EMA provides stable smoothing with constant memory:
    EMA_t = (1 - alpha) * EMA_{t-1} + alpha * x_t

    Args:
        previous (float): The EMA value from the previous step.
        new_value (float): The current observation.
        alpha (float): Smoothing factor in the range (0, 1].

    Returns:
        float: The updated EMA.
    """
    if alpha <= 0.0:
        return previous
    if alpha >= 1.0:
        return new_value
    return (1.0 - alpha) * previous + alpha * new_value


@dataclass
class ConvergenceSignals:
    """
    Container for the current optimization dynamics observables.

    Attributes:
        generation (int): Current global generation index.
        discovery_rate (float): EMA of novel structures identified.
        progress_rate (float): EMA of objective improvement events.
        acceptance_rate (float): Smoothed ratio of accepted vs proposed moves.
        stall_fraction (float): Normalized measure [0, 1] of stagnation relative to threshold.
        improvement (bool): True if an objective improvement occurred this step.
        new_structures (int): Raw count of newly discovered structures.
        accepted_structures (int): Raw count of accepted structures.
        proposed_structures (int): Raw count 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):
    """
    Engine-agnostic convergence monitor and novelty tracker.

    Tracks optimization status via objective history and optional feature-based novelty maps.
    It exposes pure "signals" consumed by thermostats or adaptive variation operators.

    Complexity:
        - track_features=False: O(1) time and space per generation.
        - track_features=True: O(N) time (relative to batch size) and O(B) space (bounded by max_cache_size).

    Args:
        stall_threshold (Optional[int]): Generations without improvement before stop.
        objective_threshold (float): Precision for defining an "improvement".
        alpha_discovery (float): Smoothing for novel structure rate.
        alpha_progress (float): Smoothing for objective progress events.
        alpha_acceptance (float): Smoothing for transition acceptance ratios.
        max_cache_size (int): Max entries in feature-objective map before eviction.
        track_features (bool): Enable per-feature objective tracking for diverse populations.
    """

    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 = 1_000_000,
        track_features: bool = False,
    ):
        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)

        # Internals
        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()
        self._best_objectives: Dict[Tuple, np.ndarray] = {}
        self._global_best = None

    def _evict_cache_if_needed(self):
        """
        Memory-bounding strategy: Purges 25% oldest entries on cache overflow.
        Uses Python dict insertion-order for ultra-fast approximate LRU.
        """
        if len(self._best_objectives) > self._max_cache_size:
            keys = list(self._best_objectives.keys())
            cutoff = int(len(keys) * 0.25)
            self._best_objectives = {k: self._best_objectives[k] for k in keys[cutoff:]}

    def check(
        self,
        generation: int,
        objectives: np.ndarray,
        features: np.ndarray,
        accepted_structures: int = 0,
        proposed_structures: int = 0,
    ) -> Dict[str, Any]:
        """
        Process a batch of results and update convergence signals.

        Args:
            generation (int): Current process generation.
            objectives (np.ndarray): Objective values [Batch, M].
            features (np.ndarray): Feature vectors [Batch, F].
            accepted_structures (int): Count of structures surviving transition.
            proposed_structures (int): Count of structures proposed to transition.

        Returns:
            Dict[str, Any]: Standardized signal dictionary for engine consumption.
        """
        self._generation = generation

        # Normalization of inputs
        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

        # --- Novelty & Progress Tracking Logic ---
        if not self._track_features:
            # Mode A: Ultra-fast Global Best tracking (O(1) memory)
            # Use nanmin to ignore NaNs if possible, fall back to INF if all are NaN
            try:
                current_min = float(np.nanmin(objectives_arr))
            except ValueError:
                current_min = float('inf')

            if np.isfinite(current_min):
                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:
            # Mode B: Per-feature Niche tracking (O(B) memory)
            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:
                    # Improvement if better in ANY objective dimension (for Pareto niching)
                    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()

            self._evict_cache_if_needed()

        self._improvement_found = improvement_found

        # Update stall logic
        if improvement_found:
            self._stall_counter = 0
        else:
            if np.isfinite(self._stall_threshold):
                self._stall_counter = min(int(self._stall_threshold), self._stall_counter + 1)
            else:
                self._stall_counter += 1

        # --- Signal Smoothing (EMA) ---
        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)

        acceptance = (accepted_structures / proposed_structures) if proposed_structures > 0 else (1.0 if accepted_structures > 0 else 0.0)
        self._ema_acceptance = _ema(self._ema_acceptance, float(acceptance), self._alpha_acceptance)

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

        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,  # Legacy support
        }

    def is_converge(self) -> bool:
        """Returns True if the termination condition (stagnation) is met."""
        return self._converged
        
    def improvement_found(self) -> bool:
        """Returns True if an improvement was found in the last generation."""
        return self._improvement_found

    def get_stagnation(self) -> int:
        """Returns the current number of generations since the last improvement."""
        return self._stall_counter

    @property
    def discovery_rate(self) -> float:
        """EMA of novel unique structures discovered per generation."""
        return self._ema_discovery

    @property
    def progress_rate(self) -> float:
        """EMA of how often the global or niche minimum is improved."""
        return self._ema_progress

    @property
    def acceptance_rate(self) -> float:
        """EMA of candidate structure acceptance ratio."""
        return self._ema_acceptance

    @property
    def stall_fraction(self) -> float:
        """Normalized progress toward convergence [0, 1]."""
        return _clip01(self._stall_counter / self._stall_threshold)

    def reset(self) -> None:
        """Resets all internal counters and smoothing signals to initial 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._best_objectives.clear()
        self._global_best = None
        self.signals = ConvergenceSignals()


# =============================================================================
# TODO & FUTURE SCALABILITY ROADMAP (CS Improvements)
# =============================================================================
# 1. SLIDING-VARIANCE DETECTION:
#    Implement objective variance monitoring within a sliding window to detect 
#    local optima plateaus before the global stall_threshold is reached.
#
# 2. HYPERVOLUME-BASED CONVERGENCE:
#    For multi-objective runs, transition from scalar-min tracking to 
#    Hypervolume (HV) monitoring to measure Pareto Front stability.
#
# 3. ADVANCED CACHE EVICTION (LRU):
#    Replace dict-insertion-order eviction with a dedicated O(1) LRU map 
#    (e.g., using collections.OrderedDict) if niche memory becomes critical.
#
# 4. ASYNC SIGNAL PROCESSING:
#    Offload check() logic to a background worker for Petascale simulations 
#    where checking billion-element niching maps might block the main Engine loop.
#
# 5. STATE PERSISTENCE:
#    Extend reset() and init to support HDF5/SQL snapshot ingestion, allowing 
#    EMA signals to survive engine restarts and maintain adaptive history.
# =============================================================================