"""
Temporal smoothing for BADAS predictions.

Provides causal/streaming-safe temporal smoothing with:
- Spike detection and dampening
- Adaptive alpha based on risk level
- Rate limiting for smooth transitions
"""

from dataclasses import dataclass
from typing import Any, Dict, List, Tuple


@dataclass
class SmoothingConfig:
    """Configuration for temporal smoothing of predictions."""

    enabled: bool = True
    max_history: int = 10
    spike_threshold: float = 0.4
    spike_baseline: float = 0.3
    spike_dampening: float = 0.15
    alpha_rise: float = 0.7
    alpha_fall: float = 0.3
    alpha_normal: float = 0.2
    max_change: float = 0.2

    @classmethod
    def default(cls) -> 'SmoothingConfig':
        return cls()

    @classmethod
    def fast_response(cls) -> 'SmoothingConfig':
        return cls(
            alpha_rise=0.9,
            alpha_fall=0.5,
            alpha_normal=0.4,
            max_change=0.3
        )

    @classmethod
    def stable(cls) -> 'SmoothingConfig':
        return cls(
            alpha_rise=0.5,
            alpha_fall=0.2,
            alpha_normal=0.1,
            max_change=0.1,
            max_history=15
        )

    @classmethod
    def disabled(cls) -> 'SmoothingConfig':
        return cls(enabled=False)


@dataclass
class PostProcessSmoothingConfig:
    """Configuration for offline post-processing smoothing."""

    enabled: bool = True
    max_history: int = 10
    spike_threshold: float = 0.4
    spike_baseline: float = 0.3
    spike_dampening: float = 0.15
    alpha_rise: float = 0.7
    alpha_fall: float = 0.3
    alpha_normal: float = 0.2
    max_change: float = 0.2

    @classmethod
    def from_smoothing_config(cls, config: SmoothingConfig) -> 'PostProcessSmoothingConfig':
        """Create from existing SmoothingConfig."""
        return cls(
            enabled=config.enabled,
            max_history=config.max_history,
            spike_threshold=config.spike_threshold,
            spike_baseline=config.spike_baseline,
            spike_dampening=config.spike_dampening,
            alpha_rise=config.alpha_rise,
            alpha_fall=config.alpha_fall,
            alpha_normal=config.alpha_normal,
            max_change=config.max_change
        )


class TemporalSmoother:
    """Temporal smoothing for predictions with spike detection and rate limiting."""

    def __init__(self, config: SmoothingConfig):
        self.config = config
        self.history: List[Tuple[int, float, float]] = []

    def reset(self):
        self.history.clear()

    def smooth(self, window_end: int, raw_prediction: float) -> float:
        if not self.config.enabled:
            self._update_history(window_end, raw_prediction, raw_prediction)
            return raw_prediction

        if not self.history:
            self._update_history(window_end, raw_prediction, raw_prediction)
            return raw_prediction

        cfg = self.config

        recent_window = min(3, len(self.history))
        recent_preds = [h[1] for h in self.history[-recent_window:]]
        recent_avg = sum(recent_preds) / len(recent_preds)

        is_spike = (
            (raw_prediction - recent_avg > cfg.spike_threshold) and
            (recent_avg < cfg.spike_baseline)
        )
        base_pred = recent_avg + cfg.spike_dampening if is_spike else raw_prediction

        prev_smooth = self.history[-1][2]

        if base_pred > prev_smooth and base_pred > 0.5:
            alpha = cfg.alpha_rise
        elif prev_smooth > 0.5 and base_pred < prev_smooth:
            alpha = cfg.alpha_fall
        else:
            alpha = cfg.alpha_normal

        smoothed = alpha * base_pred + (1 - alpha) * prev_smooth

        smoothed = max(
            prev_smooth - cfg.max_change,
            min(smoothed, prev_smooth + cfg.max_change)
        )

        self._update_history(window_end, raw_prediction, smoothed)
        return smoothed

    def _update_history(self, window_end: int, raw: float, smoothed: float):
        self.history.append((window_end, raw, smoothed))
        if len(self.history) > self.config.max_history:
            self.history.pop(0)


def apply_smoothing_offline(
    results: List[Dict[str, Any]],
    config: PostProcessSmoothingConfig,
    use_raw: bool = True,
    risk_config=None,
) -> List[Dict[str, Any]]:
    """
    Apply smoothing to results offline (post-processing).

    Args:
        results: List of prediction dictionaries from predict_video
        config: Smoothing configuration
        use_raw: If True, smooth from raw_probability. If False, smooth from probability.

    Returns:
        New list with updated probability values (original list unchanged)
    """
    # Import here to avoid circular dependency
    from badas.inference.risk import create_risk_assessment

    if not results or not config.enabled:
        return results

    # Create deep copy to avoid modifying original
    smoothed_results = []
    history: List[float] = []
    prev_smoothed: float = None

    for pred in results:
        # Get the source probability
        source_prob = pred['raw_probability'] if use_raw else pred['probability']

        # First prediction - no smoothing
        if prev_smoothed is None:
            smoothed = source_prob
        else:
            # Calculate recent average for spike detection
            recent_window = min(3, len(history))
            recent_preds = history[-recent_window:]
            recent_avg = sum(recent_preds) / len(recent_preds)

            # Spike detection
            is_spike = (
                (source_prob - recent_avg > config.spike_threshold) and
                (recent_avg < config.spike_baseline)
            )
            base_pred = recent_avg + config.spike_dampening if is_spike else source_prob

            # Adaptive alpha selection
            if base_pred > prev_smoothed and base_pred > 0.5:
                alpha = config.alpha_rise
            elif prev_smoothed > 0.5 and base_pred < prev_smoothed:
                alpha = config.alpha_fall
            else:
                alpha = config.alpha_normal

            # EMA smoothing
            smoothed = alpha * base_pred + (1 - alpha) * prev_smoothed

            # Rate limiting
            smoothed = max(
                prev_smoothed - config.max_change,
                min(smoothed, prev_smoothed + config.max_change)
            )

        # Update history
        history.append(source_prob)
        if len(history) > config.max_history:
            history.pop(0)
        prev_smoothed = smoothed

        # Create new prediction dict with updated values
        new_pred = pred.copy()
        new_pred['probability'] = round(smoothed, 4)
        risk = create_risk_assessment(smoothed, risk_config)
        new_pred['risk_level'] = risk['level']
        new_pred['risk_emoji'] = risk['emoji']

        smoothed_results.append(new_pred)

    return smoothed_results


def compare_smoothing_configs(
    results: List[Dict[str, Any]],
    configs: Dict[str, PostProcessSmoothingConfig]
) -> Dict[str, List[Dict[str, Any]]]:
    """
    Apply multiple smoothing configs and return all results for comparison.

    Args:
        results: Original prediction results
        configs: Dictionary of {name: config} pairs

    Returns:
        Dictionary of {name: smoothed_results}
    """
    return {
        name: apply_smoothing_offline(results, config)
        for name, config in configs.items()
    }


__all__ = [
    'SmoothingConfig',
    'PostProcessSmoothingConfig',
    'TemporalSmoother',
    'apply_smoothing_offline',
    'compare_smoothing_configs',
]
