#!/usr/bin/env python3
"""
SurpriseTracker - Closes the control loop for Se-aware routing.

This is the missing piece that turns the router into a controller:
1. PREDICT: Given (action, telemetry_before), predict telemetry_after
2. OBSERVE: Measure actual telemetry after action window
3. SCORE: Compute surprise = how wrong was the prediction?
4. UPDATE: Update baseline (Omega*) and learned action deltas

Per ChatGPT's recommendation, we use normalized weighted error (not KL)
because telemetry is mixed units and not clean distributions.

surprise = sum_j( w_j * |y_j - y_hat_j| / (sigma_j + epsilon) )

Where:
- y_j = observed metric
- y_hat_j = predicted metric
- sigma_j = rolling stddev (EMA) per metric
- w_j = importance weight (thermal metrics heavier)
"""

import numpy as np
from typing import Dict, Tuple, Optional, List
from dataclasses import dataclass, field
import time
import json

from .telemetry import TelemetrySnapshot


@dataclass
class MetricStats:
    """Rolling statistics for a single metric."""
    mean: float = 0.0
    var: float = 1.0  # Variance
    count: int = 0
    ema_alpha: float = 0.1

    @property
    def std(self) -> float:
        return np.sqrt(max(self.var, 1e-6))

    def update(self, value: float):
        """Update rolling mean and variance with EMA."""
        if self.count == 0:
            self.mean = value
            self.var = 1.0
        else:
            # EMA update
            delta = value - self.mean
            self.mean += self.ema_alpha * delta
            # Welford's online variance with EMA
            self.var = (1 - self.ema_alpha) * self.var + self.ema_alpha * (delta ** 2)
        self.count += 1


@dataclass
class ActionDeltas:
    """Learned per-backend telemetry deltas."""
    # Expected change in each metric when this backend is active
    gpu_temp_delta: float = 0.0
    gpu_power_delta: float = 0.0
    cpu_load_delta: float = 0.0
    mem_pressure_delta: float = 0.0
    io_wait_delta: float = 0.0
    update_count: int = 0
    ema_alpha: float = 0.15

    def update(self, observed_deltas: Dict[str, float]):
        """Update learned deltas with new observation."""
        if self.update_count == 0:
            # First observation - initialize
            self.gpu_temp_delta = observed_deltas.get('gpu_temp', 0.0)
            self.gpu_power_delta = observed_deltas.get('gpu_power', 0.0)
            self.cpu_load_delta = observed_deltas.get('cpu_load', 0.0)
            self.mem_pressure_delta = observed_deltas.get('mem_pressure', 0.0)
            self.io_wait_delta = observed_deltas.get('io_wait', 0.0)
        else:
            # EMA update
            self.gpu_temp_delta += self.ema_alpha * (observed_deltas.get('gpu_temp', 0.0) - self.gpu_temp_delta)
            self.gpu_power_delta += self.ema_alpha * (observed_deltas.get('gpu_power', 0.0) - self.gpu_power_delta)
            self.cpu_load_delta += self.ema_alpha * (observed_deltas.get('cpu_load', 0.0) - self.cpu_load_delta)
            self.mem_pressure_delta += self.ema_alpha * (observed_deltas.get('mem_pressure', 0.0) - self.mem_pressure_delta)
            self.io_wait_delta += self.ema_alpha * (observed_deltas.get('io_wait', 0.0) - self.io_wait_delta)
        self.update_count += 1

    def to_dict(self) -> Dict[str, float]:
        return {
            'gpu_temp': self.gpu_temp_delta,
            'gpu_power': self.gpu_power_delta,
            'cpu_load': self.cpu_load_delta,
            'mem_pressure': self.mem_pressure_delta,
            'io_wait': self.io_wait_delta,
        }


@dataclass
class ThermalModel:
    """
    Physics-based thermal model for GPU cooling prediction.

    Uses exponential decay toward ambient:
        T(t+Δt) = T_amb + (T_current - T_amb) * e^(-k * Δt)

    Where:
        - T_amb = ambient/idle temperature (learned from min observed)
        - k = cooling rate constant (inverse time constant)

    This fixes the "GPU cools when routed away" gap that pure delta models miss.
    """
    # Ambient temperatures (learned from minimum observed or idle periods)
    gpu_temp_ambient: float = 35.0   # Default ambient GPU temp (°C)
    gpu_power_ambient: float = 15.0  # Default idle GPU power (W)

    # Cooling rate constants (1/τ where τ is time constant in seconds)
    # Higher k = faster cooling
    # Tuned from real Quadro T2000 data: 56°C → 53°C over ~30s
    gpu_temp_k: float = 0.03   # ~33 second time constant for temp
    gpu_power_k: float = 0.3   # ~3 second time constant for power

    # Track minimum observed values to learn ambient
    gpu_temp_min: float = 100.0
    gpu_power_min: float = 500.0
    observations: int = 0
    ema_alpha: float = 0.05  # Slow EMA for ambient learning

    def cool_toward_ambient(
        self,
        current_temp: float,
        current_power: float,
        elapsed_s: float
    ) -> Tuple[float, float]:
        """
        Apply exponential cooling toward ambient.

        Returns (predicted_temp, predicted_power) after elapsed_s of cooling.
        """
        # T(t+Δt) = T_amb + (T_current - T_amb) * e^(-k * Δt)
        decay_temp = np.exp(-self.gpu_temp_k * elapsed_s)
        decay_power = np.exp(-self.gpu_power_k * elapsed_s)

        new_temp = self.gpu_temp_ambient + (current_temp - self.gpu_temp_ambient) * decay_temp
        new_power = self.gpu_power_ambient + (current_power - self.gpu_power_ambient) * decay_power

        return new_temp, new_power

    def update_ambient(self, observed_temp: float, observed_power: float, was_idle: bool = False):
        """
        Update ambient estimates from observations.

        If was_idle=True, use this observation directly.
        Otherwise, track minimum as proxy for ambient.
        """
        self.observations += 1

        if was_idle:
            # Direct observation of idle state
            self.gpu_temp_ambient += self.ema_alpha * (observed_temp - self.gpu_temp_ambient)
            self.gpu_power_ambient += self.ema_alpha * (observed_power - self.gpu_power_ambient)
        else:
            # Track minimum as proxy for ambient
            if observed_temp > 0 and observed_temp < self.gpu_temp_min:
                self.gpu_temp_min = observed_temp
                # Slowly pull ambient toward observed minimum
                self.gpu_temp_ambient += self.ema_alpha * (observed_temp - self.gpu_temp_ambient)

            if observed_power > 0 and observed_power < self.gpu_power_min:
                self.gpu_power_min = observed_power
                self.gpu_power_ambient += self.ema_alpha * (observed_power - self.gpu_power_ambient)

    def info(self) -> Dict[str, float]:
        return {
            'gpu_temp_ambient': self.gpu_temp_ambient,
            'gpu_power_ambient': self.gpu_power_ambient,
            'gpu_temp_k': self.gpu_temp_k,
            'gpu_power_k': self.gpu_power_k,
            'gpu_temp_min_observed': self.gpu_temp_min,
            'gpu_power_min_observed': self.gpu_power_min,
            'observations': self.observations,
        }


@dataclass
class SurpriseResult:
    """Result of surprise scoring."""
    total_surprise: float           # Weighted sum of z-scores (for anomaly detection)
    total_mae: float                # Weighted sum of absolute errors (for learning assessment)
    per_metric_z: Dict[str, float]  # Z-score per metric
    per_metric_error: Dict[str, float]  # Raw error per metric
    prediction: Dict[str, float]    # What we predicted
    observation: Dict[str, float]   # What we observed
    backend: str                    # Which backend was used
    se_norm: float                  # Se at decision time


class SurpriseTracker:
    """
    Tracks prediction error (surprise) to close the control loop.

    Usage:
        tracker = SurpriseTracker()

        # Before action
        telem_before = telemetry.snapshot()
        backend = qubo_router.route(...).backend

        # After action window
        telem_after = telemetry.snapshot()

        # Score and update
        pred = tracker.predict(backend, telem_before, se_norm)
        result = tracker.score(pred, telem_after, backend, se_norm)
        tracker.update(telem_after, backend, telem_before)
    """

    def __init__(self, ema_alpha: float = 0.1):
        self.ema_alpha = ema_alpha

        # Rolling stats per metric (for normalization)
        self.stats: Dict[str, MetricStats] = {
            'gpu_temp': MetricStats(ema_alpha=ema_alpha),
            'gpu_power': MetricStats(ema_alpha=ema_alpha),
            'cpu_load': MetricStats(ema_alpha=ema_alpha),
            'mem_pressure': MetricStats(ema_alpha=ema_alpha),
            'io_wait': MetricStats(ema_alpha=ema_alpha),
        }

        # Learned action deltas per backend
        self.action_deltas: Dict[str, ActionDeltas] = {
            'cpu': ActionDeltas(ema_alpha=0.15),
            'gpu': ActionDeltas(ema_alpha=0.15),
            'qpu': ActionDeltas(ema_alpha=0.15),
        }

        # Importance weights (thermal metrics heavier)
        self.weights: Dict[str, float] = {
            'gpu_temp': 2.0,    # Thermal is critical
            'gpu_power': 1.5,   # Power relates to thermal
            'cpu_load': 1.0,    # Standard
            'mem_pressure': 0.8,
            'io_wait': 0.5,
        }

        # Sigma floors per metric (prevents z-score inflation as model tightens)
        # These are minimum standard deviations - below this, we cap sigma
        self.sigma_floor: Dict[str, float] = {
            'gpu_temp': 1.0,      # 1°C minimum noise floor
            'gpu_power': 5.0,     # 5W minimum noise floor
            'cpu_load': 0.1,      # 0.1 load minimum noise floor
            'mem_pressure': 0.5,  # 0.5% minimum noise floor
            'io_wait': 0.1,       # 0.1% minimum noise floor
        }

        # Equilibrium baseline (Omega*)
        self.omega_star: Dict[str, float] = {}

        # Physics-based thermal model for GPU cooling
        self.thermal_model = ThermalModel()

        # History for analysis
        self.history: List[SurpriseResult] = []

    def _extract_metrics(self, telem: TelemetrySnapshot) -> Dict[str, float]:
        """Extract flat metric dict from telemetry snapshot."""
        return {
            'gpu_temp': telem.gpu.temp_c if telem.gpu.available else 0.0,
            'gpu_power': telem.gpu.power_w if telem.gpu.available else 0.0,
            'cpu_load': telem.cpu.load_1m,
            'mem_pressure': telem.memory.pressure_pct,
            'io_wait': telem.io.io_wait_pct,
        }

    def predict(
        self,
        backend: str,
        telem_before: TelemetrySnapshot,
        se_norm: float,
        queue_depth: int = 0,
        elapsed_s: float = 0.3
    ) -> Dict[str, float]:
        """
        Predict telemetry after action window.

        Hybrid model (per ChatGPT/WaveDynamic recommendation):
        - Physics prior: GPU cools exponentially toward ambient when not under load
        - Learned deltas: Active effects when backend is running (heating, power draw)

        Args:
            backend: Which backend will run
            telem_before: Telemetry snapshot before action
            se_norm: Symbolic entropy at decision time
            queue_depth: Number of queued tasks
            elapsed_s: Action window duration in seconds (for cooling calculation)
        """
        current = self._extract_metrics(telem_before)
        deltas = self.action_deltas.get(backend, ActionDeltas()).to_dict()

        # Queue pressure adds to CPU load and IO wait
        queue_factor = min(1.0, queue_depth / 10.0) if queue_depth > 0 else 0.0

        prediction = {}

        for metric in current:
            base = current[metric]
            delta = deltas.get(metric, 0.0)

            # GPU thermal metrics: use physics-based cooling when not using GPU
            if metric == 'gpu_temp' and backend != 'gpu' and base > 0:
                # GPU is idle - apply exponential cooling toward ambient
                pred_temp, _ = self.thermal_model.cool_toward_ambient(
                    current_temp=base,
                    current_power=current.get('gpu_power', 0.0),
                    elapsed_s=elapsed_s
                )
                prediction[metric] = pred_temp
                continue

            if metric == 'gpu_power' and backend != 'gpu' and base > 0:
                # GPU is idle - power drops toward idle
                _, pred_power = self.thermal_model.cool_toward_ambient(
                    current_temp=current.get('gpu_temp', 0.0),
                    current_power=base,
                    elapsed_s=elapsed_s
                )
                prediction[metric] = pred_power
                continue

            # For GPU backend or non-thermal metrics: use learned deltas
            # Queue pressure affects cpu_load and io_wait
            if metric == 'cpu_load':
                delta += queue_factor * 0.5
            elif metric == 'io_wait':
                delta += queue_factor * 0.2

            prediction[metric] = base + delta

        return prediction

    def score(
        self,
        prediction: Dict[str, float],
        telem_after: TelemetrySnapshot,
        backend: str,
        se_norm: float
    ) -> SurpriseResult:
        """
        Compute surprise score: how wrong was our prediction?

        Two metrics (per ChatGPT recommendation):
        - total_surprise (z-score sum): for anomaly detection
        - total_mae (weighted abs error): for learning assessment

        Sigma floor prevents z-score inflation as model tightens.
        """
        observed = self._extract_metrics(telem_after)
        epsilon = 1e-6

        per_metric_z = {}
        per_metric_error = {}
        weighted_z_sum = 0.0
        weighted_mae_sum = 0.0

        for metric in prediction:
            y_hat = prediction[metric]
            y = observed.get(metric, y_hat)

            # Apply sigma floor to prevent z-score inflation
            raw_sigma = self.stats[metric].std
            sigma_floor = self.sigma_floor.get(metric, 0.1)
            sigma = max(raw_sigma, sigma_floor)

            error = abs(y - y_hat)
            z_score = error / (sigma + epsilon)
            weight = self.weights.get(metric, 1.0)

            per_metric_error[metric] = error
            per_metric_z[metric] = z_score
            weighted_z_sum += weight * z_score
            weighted_mae_sum += weight * error

        result = SurpriseResult(
            total_surprise=weighted_z_sum,
            total_mae=weighted_mae_sum,
            per_metric_z=per_metric_z,
            per_metric_error=per_metric_error,
            prediction=prediction,
            observation=observed,
            backend=backend,
            se_norm=se_norm
        )

        self.history.append(result)
        return result

    def update(
        self,
        telem_after: TelemetrySnapshot,
        backend: str,
        telem_before: TelemetrySnapshot,
        skip_action_learning: bool = False
    ):
        """
        Update rolling stats and learned action deltas.

        Call this after scoring to improve future predictions.

        Args:
            skip_action_learning: If True, don't update action_deltas for this backend.
                                  Use for simulated backends (e.g., QPU) to prevent
                                  learning fake actuator models.
        """
        before = self._extract_metrics(telem_before)
        after = self._extract_metrics(telem_after)

        # Update rolling stats per metric (always do this)
        for metric, value in after.items():
            self.stats[metric].update(value)

        # Compute observed deltas and update backend-specific model
        # Skip for simulated backends to prevent learning fake models
        if not skip_action_learning:
            observed_deltas = {
                metric: after[metric] - before[metric]
                for metric in before
            }
            self.action_deltas[backend].update(observed_deltas)

        # Update equilibrium (Omega*) with EMA
        for metric, value in after.items():
            if metric not in self.omega_star:
                self.omega_star[metric] = value
            else:
                self.omega_star[metric] = (
                    (1 - self.ema_alpha) * self.omega_star[metric] +
                    self.ema_alpha * value
                )

        # Update thermal model ambient estimates
        # GPU was "idle" if we routed to a non-GPU backend
        gpu_was_idle = (backend != 'gpu')
        gpu_temp = after.get('gpu_temp', 0.0)
        gpu_power = after.get('gpu_power', 0.0)
        if gpu_temp > 0 or gpu_power > 0:
            self.thermal_model.update_ambient(gpu_temp, gpu_power, was_idle=gpu_was_idle)

    def get_recent_surprise(self, n: int = 10) -> float:
        """Get average surprise over last n decisions."""
        if not self.history:
            return 0.0
        recent = self.history[-n:]
        return sum(r.total_surprise for r in recent) / len(recent)

    def get_backend_surprise(self, backend: str, n: int = 20) -> float:
        """Get average surprise for a specific backend."""
        backend_history = [r for r in self.history if r.backend == backend][-n:]
        if not backend_history:
            return 0.0
        return sum(r.total_surprise for r in backend_history) / len(backend_history)

    def get_recent_mae(self, n: int = 10) -> float:
        """Get average MAE over last n decisions (for learning assessment)."""
        if not self.history:
            return 0.0
        recent = self.history[-n:]
        return sum(r.total_mae for r in recent) / len(recent)

    def get_backend_mae(self, backend: str, n: int = 20) -> float:
        """Get average MAE for a specific backend."""
        backend_history = [r for r in self.history if r.backend == backend][-n:]
        if not backend_history:
            return 0.0
        return sum(r.total_mae for r in backend_history) / len(backend_history)

    def info(self) -> Dict:
        """Return tracker state for diagnostics."""
        return {
            'history_length': len(self.history),
            'recent_surprise': self.get_recent_surprise(),
            'recent_mae': self.get_recent_mae(),
            'backend_update_counts': {
                backend: deltas.update_count
                for backend, deltas in self.action_deltas.items()
            },
            'backend_mae': {
                backend: self.get_backend_mae(backend)
                for backend in self.action_deltas.keys()
            },
            'learned_deltas': {
                backend: deltas.to_dict()
                for backend, deltas in self.action_deltas.items()
            },
            'omega_star': self.omega_star.copy(),
            'thermal_model': self.thermal_model.info(),
            'metric_stats': {
                metric: {'mean': s.mean, 'std': s.std, 'count': s.count}
                for metric, s in self.stats.items()
            },
        }

    def to_json(self) -> str:
        """Serialize tracker state to JSON."""
        return json.dumps(self.info(), indent=2)


# Module-level singleton
surprise_tracker = SurpriseTracker()


def track_surprise(
    telem_before: TelemetrySnapshot,
    telem_after: TelemetrySnapshot,
    backend: str,
    se_norm: float,
    queue_depth: int = 0
) -> SurpriseResult:
    """
    Convenience function: predict, score, and update in one call.
    """
    pred = surprise_tracker.predict(backend, telem_before, se_norm, queue_depth)
    result = surprise_tracker.score(pred, telem_after, backend, se_norm)
    surprise_tracker.update(telem_after, backend, telem_before)
    return result


if __name__ == "__main__":
    # Quick test with mock data
    from .telemetry import read_telemetry
    import time

    print("SurpriseTracker test")
    print("=" * 50)

    tracker = SurpriseTracker()
    telem1 = read_telemetry()

    # Simulate prediction for CPU backend
    pred = tracker.predict('cpu', telem1, se_norm=0.3, queue_depth=2)
    print(f"Prediction: {pred}")

    # Wait a bit
    time.sleep(1)

    # Get new telemetry
    telem2 = read_telemetry()

    # Score
    result = tracker.score(pred, telem2, 'cpu', 0.3)
    print(f"Surprise: {result.total_surprise:.4f}")
    print(f"Per-metric Z: {result.per_metric_z}")

    # Update
    tracker.update(telem2, 'cpu', telem1)
    print(f"\nTracker info: {tracker.info()}")
