"""M16 — Predictive Throttling Layer.

Battre throttLL'eM (arxiv 2408.05235). Predict next 10 steps power usage par
GPU. Throttle preemptively before thermal/power limit hit.

Cible: -30% energy/token, 0% SLO break.
"""
from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn


@dataclass
class ThrottleConfig:
    """Predictive throttling hyper-params."""

    horizon_steps: int = 10
    history_steps: int = 64
    power_limit_w: float = 700.0  # B200 default
    temp_limit_c: float = 85.0
    energy_target_pct: float = 0.7  # operate at 70% of TDP for efficiency
    slo_break_max_pct: float = 0.0  # absolute zero


class PowerPredictor(nn.Module):
    """LSTM predictor: history -> next K steps power."""

    def __init__(self, cfg: ThrottleConfig = ThrottleConfig()):
        super().__init__()
        self.cfg = cfg
        self.lstm = nn.LSTM(input_size=3, hidden_size=128, num_layers=2, batch_first=True)
        self.head = nn.Linear(128, cfg.horizon_steps)

    def forward(self, history: torch.Tensor) -> torch.Tensor:
        """history: [B, T, 3] (power, util, temp) -> [B, horizon_steps] predicted power."""
        out, _ = self.lstm(history)
        last_h = out[:, -1, :]  # [B, 128]
        future_power = self.head(last_h)
        return future_power


class PredictiveThrottler:
    """Per-GPU throttling controller.

    Maintains rolling buffer of telemetry. Each forward step:
        1. Update history
        2. Predict next horizon
        3. If predicted exceeds power_limit OR temp_limit: throttle
        4. If predicted under energy_target: full speed
    """

    def __init__(self, predictor: PowerPredictor, cfg: ThrottleConfig = ThrottleConfig()):
        self.predictor = predictor
        self.cfg = cfg
        self.history: deque = deque(maxlen=cfg.history_steps)
        self.stats = {"throttles": 0, "max_predicted_w": 0.0, "energy_saved_pct": 0.0}

    def step(self, power_w: float, util_pct: float, temp_c: float) -> float:
        """Process telemetry; return throttle ratio [0.5, 1.0]."""
        self.history.append([power_w, util_pct, temp_c])
        if len(self.history) < 8:
            return 1.0  # not enough history

        hist = torch.tensor([list(self.history)], dtype=torch.float32)
        with torch.no_grad():
            future = self.predictor(hist).squeeze(0)
        max_predicted = float(future.max().item())
        self.stats["max_predicted_w"] = max(self.stats["max_predicted_w"], max_predicted)

        # Decide throttle
        if max_predicted > self.cfg.power_limit_w:
            throttle = self.cfg.power_limit_w / max_predicted
            throttle = max(0.5, throttle)  # min 50% speed
            self.stats["throttles"] += 1
            return throttle
        if temp_c > self.cfg.temp_limit_c:
            self.stats["throttles"] += 1
            return 0.7  # aggressive throttle on overheat
        return 1.0  # full speed

    def report(self) -> dict:
        return dict(self.stats)


if __name__ == "__main__":
    cfg = ThrottleConfig()
    pred = PowerPredictor(cfg)
    fake_hist = torch.randn(1, 32, 3) * 100 + 400  # ~400W base
    out = pred(fake_hist)
    print(f"[M16 throttle] predicted shape {out.shape}, sample: {out[0, :3].tolist()}")
    thr = PredictiveThrottler(pred, cfg)
    for step in range(40):
        ratio = thr.step(power_w=500.0 + step * 5, util_pct=80.0, temp_c=70.0)
    print(f"[M16 throttle] stats after 40 steps: {thr.report()}, last ratio={ratio:.3f}")
    print(f"[M16 throttle] params: {sum(p.numel() for p in pred.parameters())/1e6:.2f}M")
    print("[smoke M16] OK")