from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np

from .ecg import FloatArray, ecg_pqrst, normalize_01

Mode = Literal["normal", "run", "die", "aed"]


@dataclass(frozen=True)
class SandevistanSchedule:
    lr_min: float
    lr_max: float
    cycle_steps: int
    mode: Mode = "normal"
    aed_spikes: int = 6
    aed_spike_width: int = 1

    def __post_init__(self) -> None:
        if self.lr_min <= 0 or self.lr_max <= 0:
            raise ValueError("lr_min and lr_max must be > 0")
        if self.lr_max <= self.lr_min:
            raise ValueError("lr_max must be > lr_min")
        if self.cycle_steps <= 1:
            raise ValueError("cycle_steps must be > 1")
        if self.aed_spikes <= 0:
            raise ValueError("aed_spikes must be > 0")
        if self.aed_spike_width <= 0:
            raise ValueError("aed_spike_width must be > 0")

    def _base_wave(self, pct: FloatArray) -> FloatArray:
        raw = ecg_pqrst(pct)
        return normalize_01(raw)

    def _apply_mode(self, wave01: FloatArray) -> FloatArray:
        if self.mode == "normal":
            return wave01
        if self.mode == "run":
            return np.clip(wave01**0.85, 0.0, 1.0)
        if self.mode == "die":
            return np.zeros_like(wave01)
        if self.mode == "aed":
            out = np.zeros_like(wave01)
            n = out.shape[0]
            pos = np.linspace(0, n - 1, self.aed_spikes, dtype=int)
            half = self.aed_spike_width // 2
            for p in pos:
                lo = max(0, p - half)
                hi = min(n, p - half + self.aed_spike_width)
                out[lo:hi] = 1.0
            return out
        return wave01

    def values(self, n_steps: int) -> FloatArray:
        if n_steps <= 0:
            raise ValueError("n_steps must be > 0")
        pct = np.linspace(0.0, 1.0, n_steps, dtype=np.float64)
        wave01 = self._apply_mode(self._base_wave(pct))
        return self.lr_min + (self.lr_max - self.lr_min) * wave01

    def value(self, step: int) -> float:
        s = step % self.cycle_steps
        pct = np.array([s / (self.cycle_steps - 1)], dtype=np.float64)
        wave01 = float(self._apply_mode(self._base_wave(pct))[0])
        return float(self.lr_min + (self.lr_max - self.lr_min) * wave01)
