from __future__ import annotations

from typing import List, Tuple
import numpy as np

from .ecg_canonical import ecg_mult


class SandevistanScheduleBase:
    def __call__(self, step: int) -> float:
        return self.lr_at_step(step)

    """
    Framework-agnostic core schedule for SandevistanLR.

    This class implements the canonical ECG-shaped learning-rate logic
    and exposes a single method: `lr_at_step(step) -> float`.

    It is intentionally NumPy-only and stateless per call.
    """

    def __init__(
        self,
        *,
        steps_per_epoch: int,
        epoch_plan: List[Tuple[str, int]],
        lr_min: float,
        lr_max: float,
        die_level: float = 0.08,
        aed_base: float = 0.05,
        shock_strength: float = 0.95,
        shock_sig: float = 0.010,
        shock_positions: Tuple[float, ...] = (0.30, 0.55, 0.80),
        norm_samples: int = 4096,
    ) -> None:
        if steps_per_epoch <= 0:
            raise ValueError("steps_per_epoch must be > 0")

        if not epoch_plan:
            raise ValueError("epoch_plan must be non-empty")

        self.steps_per_epoch: int = int(steps_per_epoch)
        self.epoch_plan: List[Tuple[str, int]] = [
            (str(m).lower(), int(r)) for m, r in epoch_plan
        ]

        self.lr_min: float = float(lr_min)
        self.lr_max: float = float(lr_max)

        self.die_level: float = float(die_level)
        self.aed_base: float = float(aed_base)
        self.shock_strength: float = float(shock_strength)
        self.shock_sig: float = float(shock_sig)
        self.shock_positions: Tuple[float, ...] = tuple(shock_positions)

        self.norm_samples: int = int(norm_samples)

        # Precompute normalization bounds
        self._normal_min, self._normal_max = self._compute_minmax(profile="normal")
        self._run_min, self._run_max = self._compute_minmax(profile="run")

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def lr_at_step(self, step: int) -> float:
        """
        Compute learning rate at global step index.
        """
        if step < 0:
            raise ValueError("step must be >= 0")

        epoch = step // self.steps_per_epoch
        step_in_epoch = step % self.steps_per_epoch

        idx = min(epoch, len(self.epoch_plan) - 1)
        mode, repeat = self.epoch_plan[idx]

        pct = step_in_epoch / max(self.steps_per_epoch - 1, 1)

        # DIE mode (flatline)
        lr_die = self.lr_min + self.die_level * (self.lr_max - self.lr_min)
        lr_die = float(np.clip(lr_die, self.lr_min, self.lr_max))

        if mode == "die" or repeat == 0:
            return lr_die

        # AED mode (shocks)
        if mode == "aed":
            return self._aed_lr(pct, shocks_count=repeat)

        # NORMAL / RUN (ECG)
        beats = max(repeat, 1)
        phase = (pct * beats) % 1.0

        profile = "run" if mode == "run" else "normal"
        return self._ecg_lr(phase, profile=profile)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _ecg_lr(self, phase: float, *, profile: str) -> float:
        x = np.array([phase], dtype=np.float32)

        mult = ecg_mult(x)
        m = float(mult[0])

        if profile == "run":
            w = (m - self._run_min) / (self._run_max - self._run_min + 1e-12)
        else:
            w = (m - self._normal_min) / (self._normal_max - self._normal_min + 1e-12)

        w = float(np.clip(w, 0.0, 1.0))
        lr = self.lr_min + (self.lr_max - self.lr_min) * w
        return float(np.clip(lr, self.lr_min, self.lr_max))

    def _aed_lr(self, pct: float, shocks_count: int) -> float:
        base = self.lr_min + self.aed_base * (self.lr_max - self.lr_min)

        shock_vals = []
        for pos in self.shock_positions[: max(shocks_count, 0)]:
            shock_vals.append(np.exp(-0.5 * ((pct - pos) / self.shock_sig) ** 2))

        shock = max(shock_vals) if shock_vals else 0.0
        lr = base + self.shock_strength * (self.lr_max - base) * shock

        return float(np.clip(lr, self.lr_min, self.lr_max))

    def _compute_minmax(self, *, profile: str) -> Tuple[float, float]:
        xs = np.linspace(0.0, 1.0, self.norm_samples, dtype=np.float32)
        mult = ecg_mult(xs)
        return float(mult.min()), float(mult.max())


# Backward compatibility (remove in v1.0)
SandevistanSchedule = SandevistanScheduleBase
