"""
kiri/optim/schedulers.py

Learning rate schedulers. Work with any Kiri optimizer.

Usage
-----
optimizer = kiri.optim.Adam(model.parameters(), lr=0.1)
scheduler = kiri.optim.StepLR(optimizer, step_size=10, gamma=0.5)

for epoch in range(epochs):
    train(...)
    scheduler.step()
    print(scheduler.get_lr())   # see current lr
"""

import math
import numpy as np


class _Scheduler:
    def __init__(self, optimizer, last_epoch=-1):
        self.optimizer  = optimizer
        self.last_epoch = last_epoch
        self.base_lr    = optimizer.lr
        self._step_count = 0

    def step(self):
        self.last_epoch += 1
        self._step_count += 1
        new_lr = self._get_lr()
        self.optimizer.lr = new_lr
        # update MLX optimizer if present
        if hasattr(self.optimizer, '_mlx_opt') and self.optimizer._mlx_opt is not None:
            self.optimizer._mlx_opt.learning_rate = new_lr

    def get_lr(self):
        return self.optimizer.lr

    def _get_lr(self):
        raise NotImplementedError


class StepLR(_Scheduler):
    """
    Multiply LR by `gamma` every `step_size` epochs.

    lr = base_lr * gamma^(epoch // step_size)
    """
    def __init__(self, optimizer, step_size, gamma=0.1, last_epoch=-1):
        self.step_size = step_size
        self.gamma     = gamma
        super().__init__(optimizer, last_epoch)

    def _get_lr(self):
        return self.base_lr * (self.gamma ** (self.last_epoch // self.step_size))

    def __repr__(self):
        return f"StepLR(step_size={self.step_size}, gamma={self.gamma})"


class MultiStepLR(_Scheduler):
    """Decay LR by gamma at specific milestones."""
    def __init__(self, optimizer, milestones, gamma=0.1, last_epoch=-1):
        self.milestones = sorted(milestones)
        self.gamma      = gamma
        super().__init__(optimizer, last_epoch)

    def _get_lr(self):
        n_decays = sum(1 for m in self.milestones if self.last_epoch >= m)
        return self.base_lr * (self.gamma ** n_decays)


class ExponentialLR(_Scheduler):
    """Multiply LR by gamma every epoch."""
    def __init__(self, optimizer, gamma, last_epoch=-1):
        self.gamma = gamma
        super().__init__(optimizer, last_epoch)

    def _get_lr(self):
        return self.base_lr * (self.gamma ** self.last_epoch)


class CosineAnnealingLR(_Scheduler):
    """
    Cosine annealing between base_lr and eta_min over T_max epochs.

    lr = eta_min + 0.5*(base_lr - eta_min)*(1 + cos(pi * t/T_max))
    """
    def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1):
        self.T_max   = T_max
        self.eta_min = eta_min
        super().__init__(optimizer, last_epoch)

    def _get_lr(self):
        t = self.last_epoch % self.T_max
        return self.eta_min + 0.5 * (self.base_lr - self.eta_min) * (
            1 + math.cos(math.pi * t / self.T_max)
        )


class CosineAnnealingWarmRestarts(_Scheduler):
    """Cosine annealing with periodic warm restarts (SGDR)."""
    def __init__(self, optimizer, T_0, T_mult=1, eta_min=0, last_epoch=-1):
        self.T_0    = T_0
        self.T_mult = T_mult
        self.eta_min = eta_min
        super().__init__(optimizer, last_epoch)

    def _get_lr(self):
        T_cur = self.last_epoch
        T_i   = self.T_0
        while T_cur >= T_i:
            T_cur -= T_i
            T_i   *= self.T_mult
        return self.eta_min + 0.5 * (self.base_lr - self.eta_min) * (
            1 + math.cos(math.pi * T_cur / T_i)
        )


class LinearWarmup(_Scheduler):
    """
    Linearly ramps LR from 0 to base_lr over `warmup_steps` epochs,
    then holds constant. Useful before handing off to another scheduler.
    """
    def __init__(self, optimizer, warmup_epochs, last_epoch=-1):
        self.warmup_epochs = warmup_epochs
        super().__init__(optimizer, last_epoch)

    def _get_lr(self):
        if self.last_epoch < self.warmup_epochs:
            return self.base_lr * (self.last_epoch + 1) / self.warmup_epochs
        return self.base_lr


class ReduceLROnPlateau:
    """
    Reduce LR when a metric stops improving.
    Unlike other schedulers, call scheduler.step(metric_value).

    Parameters
    ----------
    mode     : 'min' (loss) or 'max' (accuracy)
    factor   : multiply LR by this when plateauing
    patience : epochs to wait before reducing
    min_lr   : floor on LR
    """
    def __init__(self, optimizer, mode='min', factor=0.5,
                 patience=5, min_lr=1e-7, verbose=True):
        self.optimizer = optimizer
        self.mode      = mode
        self.factor    = factor
        self.patience  = patience
        self.min_lr    = min_lr
        self.verbose   = verbose
        self._best     = float('inf') if mode == 'min' else float('-inf')
        self._wait     = 0

    def step(self, metric):
        improved = (metric < self._best) if self.mode == 'min' else (metric > self._best)
        if improved:
            self._best = metric
            self._wait = 0
        else:
            self._wait += 1
            if self._wait >= self.patience:
                old_lr = self.optimizer.lr
                new_lr = max(old_lr * self.factor, self.min_lr)
                self.optimizer.lr = new_lr
                if hasattr(self.optimizer, '_mlx_opt') and self.optimizer._mlx_opt is not None:
                    self.optimizer._mlx_opt.learning_rate = new_lr
                if self.verbose and new_lr != old_lr:
                    print(f"  ReduceLROnPlateau: lr {old_lr:.2e} → {new_lr:.2e}")
                self._wait = 0

    def get_lr(self):
        return self.optimizer.lr
