"""
Convergio — Automatic Convergence Detection for Iterative Methods
================================================================
Detects convergence, oscillation, bistability, and divergence
in any numerical time series. One function call. No configuration needed.

Usage:
    from convergio import detect, watch

    # Analyze a completed run
    result = detect(my_time_series)
    print(result.state)        # "converged" | "oscillating" | "bistable" | "diverging"
    print(result.converged_at) # step where convergence was detected
    print(result.quality)      # 0.0 - 1.0

    # Monitor a running simulation
    mon = watch(var_threshold=1e-6)
    for step in range(10000):
        value = my_solver_step()
        should_stop, info = mon.step(value)
        if should_stop:
            print(f"Converged at step {step}, saved {10000-step} steps")
            break

(c) 2026 Maximilian Jurak — 3 Kaiserberge Engineering & Research
"""

__version__ = "0.1.4"
__author__ = "Maximilian Jurak"

import numpy as np
from typing import Optional, Tuple, List, Union
from dataclasses import dataclass, field


# =============================================================
# Result Object
# =============================================================

@dataclass
class ConvergenceResult:
    """Result of convergence analysis."""

    state: str
    """One of: 'converged', 'stalling', 'oscillating', 'bistable', 'diverging', 'warming_up', 'unconverged'"""

    quality: float
    """Convergence quality score from 0.0 (no convergence) to 1.0 (perfect)."""

    converged_at: Optional[int]
    """Step index where convergence was first detected. None if not converged."""

    final_variance: float
    """Variance of the signal in the last analysis window."""

    oscillation_freq: float
    """Detected oscillation frequency (0.0 if not oscillating)."""

    n_modes: int
    """Number of distinct stable modes detected (1 = unimodal, 2+ = multimodal)."""

    recommendation: str
    """One of: 'stop', 'continue', 'restart'."""

    saved_steps: int = 0
    """Steps saved by early stopping (0 if not applicable)."""

    total_steps: int = 0
    """Total steps in the analyzed series."""

    def __repr__(self):
        return (f"ConvergenceResult(state='{self.state}', quality={self.quality:.2f}, "
                f"converged_at={self.converged_at}, recommendation='{self.recommendation}')")


# =============================================================
# Core Analysis
# =============================================================

def rolling_variance(signal: np.ndarray, window: int = 50) -> np.ndarray:
    """
    Compute rolling variance of a signal.

    Args:
        signal: 1D array of values
        window: window size for variance computation

    Returns:
        Array of rolling variance values
    """
    signal = np.asarray(signal, dtype=float)
    n = len(signal)

    if n < window:
        return np.array([np.var(signal)])

    # Efficient rolling variance using cumulative sums
    result = np.zeros(n - window + 1)
    for i in range(len(result)):
        result[i] = np.var(signal[i:i + window])

    return result


def count_modes(signal: np.ndarray, n_bins: int = 30) -> int:
    """
    Count distinct modes (peaks in distribution) of signal values.

    Args:
        signal: 1D array of values
        n_bins: number of histogram bins

    Returns:
        Number of detected modes (>= 1)
    """
    signal = np.asarray(signal, dtype=float)

    if len(signal) < 10:
        return 1

    hist, _ = np.histogram(signal, bins=n_bins)
    hist_smooth = np.convolve(hist, [0.25, 0.5, 0.25], mode='same')

    peaks = 0
    threshold = np.max(hist_smooth) * 0.15

    for i in range(1, len(hist_smooth) - 1):
        if (hist_smooth[i] > hist_smooth[i - 1] and
            hist_smooth[i] > hist_smooth[i + 1] and
            hist_smooth[i] > threshold):
            peaks += 1

    return max(1, peaks)


def detect(
    signal: Union[list, np.ndarray],
    window: int = 0,
    var_threshold: float = 1e-6,
    osc_threshold: float = 0.3,
    residual_target: float = 0.0,
) -> ConvergenceResult:
    """
    Analyze a time series for convergence.

    This is the main entry point. Pass any 1D numerical time series
    and get back a complete convergence analysis.

    Args:
        signal: 1D array of scalar values (residuals, energies, loss, etc.)
        window: analysis window size (0 = auto, typically 5-10% of signal length)
        var_threshold: variance below this = signal is stable
        osc_threshold: zero-crossing frequency above this = oscillating
        residual_target: if > 0, signal must be BELOW this to count as "converged".
                         A stable signal ABOVE this target = "stalling" (plateau).

    Returns:
        ConvergenceResult with state, quality, recommendation, etc.

    Examples:
        >>> import numpy as np
        >>> from convergio import detect
        >>>
        >>> # Converging signal
        >>> signal = np.exp(-np.linspace(0, 5, 1000)) + np.random.randn(1000) * 0.001
        >>> result = detect(signal)
        >>> result.state
        'converged'
        >>>
        >>> # Stalling signal (stable but above target)
        >>> signal = np.ones(1000) * 0.01 + np.random.randn(1000) * 0.0001
        >>> result = detect(signal, residual_target=1e-4)
        >>> result.state
        'stalling'
    """
    signal = np.asarray(signal, dtype=float)
    n = len(signal)

    # Auto window size
    if window <= 0:
        window = max(10, min(200, n // 10))

    # Not enough data
    if n < window * 2:
        return ConvergenceResult(
            state="warming_up",
            quality=0.0,
            converged_at=None,
            final_variance=float(np.var(signal)) if n > 1 else 0.0,
            oscillation_freq=0.0,
            n_modes=1,
            recommendation="continue",
            total_steps=n,
        )

    # --- Rolling variance ---
    rv = rolling_variance(signal, window)
    final_var = float(rv[-1]) if len(rv) > 0 else float(np.var(signal))

    # --- Check divergence (variance growing or signal unbounded) ---
    if len(rv) > 10:
        tail = rv[-min(20, len(rv)):]
        trend = np.polyfit(np.arange(len(tail)), tail, 1)[0]

        # Compare first and last quarter variance
        q1 = np.mean(rv[:max(1, len(rv) // 4)])
        q4 = np.mean(rv[-max(1, len(rv) // 4):])

        # Also check if signal magnitude itself is growing (random walk)
        first_quarter_range = np.ptp(signal[:n // 4]) if n > 4 else 0
        last_quarter_range = np.ptp(signal[-n // 4:]) if n > 4 else 0

        is_growing = (
            (trend > 0 and q4 > q1 * 2) or
            (last_quarter_range > first_quarter_range * 2 and final_var > var_threshold * 10) or
            (final_var > var_threshold * 1000 and trend > 0)
        )

        if is_growing and final_var > var_threshold * 10:
            return ConvergenceResult(
                state="diverging",
                quality=0.0,
                converged_at=None,
                final_variance=final_var,
                oscillation_freq=0.0,
                n_modes=1,
                recommendation="restart",
                total_steps=n,
            )

    # --- Find convergence point ---
    converged_at = None
    for i, v in enumerate(rv):
        if v < var_threshold:
            converged_at = i + window
            break

    # --- Oscillation detection ---
    detrended = signal - np.mean(signal)
    zero_crossings = int(np.sum(np.diff(np.sign(detrended)) != 0))
    osc_freq = zero_crossings / (2.0 * n)

    # Check if variance is stable (not decreasing) — sign of sustained oscillation
    variance_stable = False
    if len(rv) > 4:
        first_half_var = np.mean(rv[:len(rv) // 2])
        second_half_var = np.mean(rv[len(rv) // 2:])
        if first_half_var > 0 and second_half_var / (first_half_var + 1e-30) > 0.5:
            variance_stable = True

    # Oscillating = enough zero crossings + variance not decreasing + not converged
    is_oscillating = (zero_crossings > n * 0.01 and
                      variance_stable and
                      final_var > var_threshold)

    # --- Mode count (bistability) ---
    tail_len = min(n, window * 4)
    n_modes = count_modes(signal[-tail_len:])

    # --- Classify ---
    # Compute tail mean for stalling detection
    tail_mean = float(np.mean(signal[-window:]))

    if final_var < var_threshold:
        # Signal is STABLE — check if still declining or truly flat
        half = window // 2
        first_half_mean = float(np.mean(signal[-(window):-(half)])) if half > 0 else tail_mean
        second_half_mean = float(np.mean(signal[-(half):])) if half > 0 else tail_mean
        still_declining = (first_half_mean > 1e-30 and
                           second_half_mean < first_half_mean * 0.95)

        below_target = (residual_target <= 0.0 or tail_mean <= residual_target)

        if below_target:
            # Below target (or no target set) → converged
            state = "converged"
            quality = min(1.0, 1.0 - final_var / var_threshold)
            recommendation = "stop"
        elif still_declining:
            # Above target but still making progress
            state = "warming_up"
            quality = 0.5
            recommendation = "continue"
        else:
            # Flat and above target → truly stalled
            state = "stalling"
            quality = 0.1
            recommendation = "adjust_params"
    elif n_modes >= 2 and not is_oscillating:
        state = "bistable"
        quality = 0.3
        recommendation = "continue"
    elif is_oscillating:
        state = "oscillating"
        quality = max(0.1, 0.5 * (1.0 - min(1.0, final_var / (var_threshold * 100))))
        recommendation = "continue"
    else:
        # Slow convergence or unclear
        if converged_at is not None:
            state = "converged"
            quality = min(1.0, 1.0 - final_var / var_threshold)
            recommendation = "stop"
        else:
            state = "unconverged"
            quality = max(0.0, 1.0 - min(1.0, final_var / (var_threshold * 1000)))
            recommendation = "continue"

    return ConvergenceResult(
        state=state,
        quality=quality,
        converged_at=converged_at,
        final_variance=final_var,
        oscillation_freq=osc_freq,
        n_modes=n_modes,
        recommendation=recommendation,
        total_steps=n,
    )


# =============================================================
# Live Monitor (for running simulations)
# =============================================================

class Watch:
    """
    Live convergence monitor for running simulations.

    Call .step(value) every iteration. Returns (should_stop, result).

    Example:
        mon = Watch(var_threshold=1e-6, check_every=100, min_steps=200)
        for i in range(10000):
            residual = solver.step()
            stop, info = mon.step(residual)
            if stop:
                print(f"Converged at step {i}")
                break
    """

    def __init__(
        self,
        var_threshold: float = 1e-6,
        check_every: int = 50,
        min_steps: int = 100,
        window: int = 0,
        residual_target: float = 0.0,
    ):
        self.var_threshold = var_threshold
        self.check_every = check_every
        self.min_steps = min_steps
        self.window = window
        self.residual_target = residual_target

        self._history: List[float] = []
        self._step_count = 0
        self._converged = False
        self._last_result: Optional[ConvergenceResult] = None

    def step(self, value: float) -> Tuple[bool, Optional[ConvergenceResult]]:
        """
        Feed one value. Returns (should_stop, result_or_None).

        Args:
            value: current scalar value (residual, energy, loss, etc.)

        Returns:
            (True, ConvergenceResult) if should stop
            (False, None) if should continue
            (False, ConvergenceResult) at check points
        """
        self._history.append(float(value))
        self._step_count += 1

        if self._converged:
            return True, self._last_result

        if self._step_count < self.min_steps:
            return False, None

        if self._step_count % self.check_every != 0:
            return False, None

        # Analyze
        result = detect(
            np.array(self._history),
            window=self.window,
            var_threshold=self.var_threshold,
            residual_target=self.residual_target,
        )
        result.total_steps = self._step_count

        self._last_result = result

        if result.recommendation == "stop" or result.state == "diverging":
            self._converged = True
            result.converged_at = self._step_count
            return True, result

        # Stalling: report but don't auto-stop (user decides)
        if result.state == "stalling":
            return False, result

        return False, result

    @property
    def history(self) -> np.ndarray:
        """Return the full recorded history."""
        return np.array(self._history)

    @property
    def n_steps(self) -> int:
        """Number of steps recorded."""
        return self._step_count

    def reset(self):
        """Reset the monitor for a new run."""
        self._history = []
        self._step_count = 0
        self._converged = False
        self._last_result = None


def watch(
    var_threshold: float = 1e-6,
    check_every: int = 50,
    min_steps: int = 100,
    window: int = 0,
    residual_target: float = 0.0,
) -> Watch:
    """
    Create a live convergence monitor.

    Args:
        var_threshold: variance below this = signal is stable
        check_every: check convergence every N steps
        min_steps: minimum steps before first check
        window: analysis window (0 = auto)
        residual_target: if > 0, signal must be BELOW this to count as converged.
                         Stable signal above target = stalling.

    Returns:
        Watch instance. Call .step(value) in your loop.
    """
    return Watch(
        var_threshold=var_threshold,
        check_every=check_every,
        min_steps=min_steps,
        window=window,
        residual_target=residual_target,
    )
