from __future__ import annotations

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

Mode = str  # "normal" | "run" | "die" | "aed"
EpochPlan = List[Tuple[Mode, int]]


@dataclass(frozen=True)
class ConfigSpec:
    epoch_plan: EpochPlan
    kwargs: Dict[str, float] | None = None
    override_lr: Dict[str, float] | None = None


CONFIGS: Dict[str, ConfigSpec] = {
    "01_default_normal_x1": ConfigSpec(
        epoch_plan=[("normal", 1)],
    ),
    "02_warmup_normal_cooldown": ConfigSpec(
        epoch_plan=([("normal", 1)] * 3)
        + ([("normal", 1)] * 14)
        + ([("normal", 1)] * 3),
    ),
    "03_athlete_training": ConfigSpec(
        epoch_plan=(
            [("normal", 1)] * 2
            + [("run", 4)] * 2
            + [("normal", 1)] * 2
            + [("run", 5)] * 2
            + [("normal", 1)] * 2
        ),
        kwargs={
            "run_t_mu_shift": -0.08,
            "run_t_sig_mul": 0.75,
            "run_r_amp_mul": 1.10,
            "run_t_amp_mul": 1.15,
            "run_gain_mul": 1.08,
        },
    ),
    "04_intermittent_flatlines": ConfigSpec(
        epoch_plan=[
            ("normal", 1),
            ("normal", 1),
            ("normal", 0),
            ("normal", 1),
            ("normal", 0),
            ("normal", 1),
            ("normal", 1),
            ("normal", 1),
            ("normal", 1),
            ("normal", 1),
        ],
    ),
    "05_die_aed_recovery": ConfigSpec(
        epoch_plan=([("normal", 1)] * 3)
        + ([("die", 0)] * 1)
        + ([("aed", 2)] * 1)
        + ([("normal", 1)] * 6),
    ),
    "06_multi_shock_protocol": ConfigSpec(
        epoch_plan=([("die", 0)] * 1)
        + ([("aed", 1)] * 1)
        + ([("aed", 2)] * 1)
        + ([("aed", 3)] * 1)
        + ([("normal", 1)] * 6),
        kwargs={
            "shock_strength": 0.98,
            "shock_sig": 0.008,
        },
    ),
    "07_interval_training": ConfigSpec(
        epoch_plan=(
            [
                ("normal", 1),
                ("run", 4),
                ("normal", 1),
                ("run", 5),
            ]
            * 3
        )
        + [("normal", 1), ("normal", 1)],
    ),
    "08_arrhythmia_stress_test": ConfigSpec(
        epoch_plan=[
            ("normal", 1),
            ("run", 6),
            ("normal", 2),
            ("normal", 0),
            ("run", 4),
            ("die", 0),
            ("aed", 1),
            ("normal", 1),
            ("run", 5),
            ("normal", 1),
        ],
    ),
    "09_conservative_small_range": ConfigSpec(
        epoch_plan=[("normal", 1)],
        override_lr={"lr_min": 5e-7, "lr_max": 5e-5},
    ),
    "10_high_energy_run_mode": ConfigSpec(
        epoch_plan=([("normal", 1)] * 2)
        + ([("run", 8)] * 1)
        + ([("normal", 1)] * 2)
        + ([("run", 10)] * 1)
        + ([("normal", 1)] * 2),
        kwargs={
            "run_t_mu_shift": -0.08,
            "run_t_sig_mul": 0.75,
            "run_r_amp_mul": 1.12,
            "run_t_amp_mul": 1.18,
            "run_gain_mul": 1.10,
        },
    ),
}
