"""Generators that provide different rates, schedules, decays or series."""
import itertools
from dataclasses import dataclass
from typing import Any, Callable, Dict, Generator, Generic, Optional, Tuple, TypeVar

import numpy

from .config import registry

OutT = TypeVar("OutT")


class Schedule(Generic[OutT]):
    """Class for implementing Thinc schedules."""

    name: str
    _schedule: Callable
    _attrs: Dict[str, Any]

    __slots__ = ["name", "_schedule", "_attrs"]

    def __init__(
        self, name: str, schedule: Callable, *, attrs: Dict[str, Any] = {}
    ) -> None:
        """Initialize a new schedule.

        name (str): The name of the schedule type.
        schedule (Callable): The schedule function.
        """
        self.name = name
        self._schedule = schedule
        self._attrs = dict(attrs)

    def __call__(self, step: int, **extra) -> OutT:
        """Compute the schedule for a given step."""

        if step < 0:
            raise ValueError(f"Step must be non-negative, was: {step}")

        return self._schedule(self, step, **extra)

    @property
    def attrs(self):
        """Schedule attributes."""
        return self._attrs

    def to_generator(
        self, start: int = 0, step_size=1, **extra
    ) -> Generator[OutT, None, None]:
        """Turn the schedule into a generator.

        start (int): The schedule initial step.
        step_size (int): The amount to increase the step for each generated value.
        **extra: Additional arguments that are passed to the schedule.
        RETURNS (Generator[OutT, None, None]): The generator.
        """
        if start < 0:
            raise ValueError(f"Schedule start must be non-negative, was: {start}")
        if step_size < 0:
            raise ValueError(f"Step size must be non-negative, was: {step_size}")

        def generate():
            for step in itertools.count(start, step_size):
                yield self(step, **extra)

        return generate()


@registry.schedules("constant_then.v1")
def constant_then(rate: OutT, steps: int, schedule: Schedule[OutT]) -> Schedule[OutT]:
    """Yield a constant rate for N steps, before starting a schedule."""
    return Schedule(
        "constant_then",
        _constant_then_schedule,
        attrs={"rate": rate, "steps": steps, "schedule": schedule},
    )


def _constant_then_schedule(schedule: Schedule, step: int, **kwargs) -> float:
    rate = schedule.attrs["rate"]
    steps = schedule.attrs["steps"]
    schedule = schedule.attrs["schedule"]

    if step < steps:
        return rate
    else:
        return schedule(step=step, **kwargs)


@registry.schedules("constant.v1")
def constant(rate: OutT) -> Schedule[OutT]:
    """Yield a constant rate."""
    return Schedule("constant", _constant_schedule, attrs={"rate": rate})


def _constant_schedule(schedule: Schedule, step: int, **kwargs) -> float:
    rate = schedule.attrs["rate"]
    return rate


@registry.schedules("decaying.v1")
def decaying(base_rate: float, decay: float, *, t: float = 0.0) -> Schedule[float]:
    """Yield an infinite series of linearly decaying values,
    following the schedule: base_rate * 1 / (1 + decay * (t + step))

    EXAMPLE:
        >>> learn_rates = decaying(0.001, 1e-4)
        >>> next(learn_rates)
        0.001
        >>> next(learn_rates)
        0.00999
    """
    return Schedule(
        "decaying",
        _decaying_schedule,
        attrs={"base_rate": base_rate, "decay": decay, "t": t},
    )


def _decaying_schedule(schedule: Schedule, step: int, **kwargs) -> float:
    base_rate = schedule.attrs["base_rate"]
    decay = schedule.attrs["decay"]
    t = schedule.attrs["t"]
    return base_rate * (1.0 / (1.0 + decay * (step + t)))


@registry.schedules("compounding.v1")
def compounding(
    start: float, stop: float, compound: float, *, t: float = 0.0
) -> Schedule[float]:
    """Yield an infinite series of compounding values. Each time the
    generator is called, a value is produced by multiplying the previous
    value by the compound rate.

    EXAMPLE:
        >>> sizes = compounding(1.0, 10.0, 1.5)
        >>> assert next(sizes) == 1.
        >>> assert next(sizes) == 1 * 1.5
        >>> assert next(sizes) == 1.5 * 1.5
    """
    return Schedule(
        "compounding",
        _compounding_schedule,
        attrs={"start": start, "stop": stop, "compound": compound, "t": t},
    )


def _compounding_schedule(schedule: Schedule, step: int, **kwargs) -> float:
    start = schedule.attrs["start"]
    stop = schedule.attrs["stop"]
    compound = schedule.attrs["compound"]
    t = schedule.attrs["t"]
    return _clip(start * (compound ** (step + t)), start, stop)


def _clip(value: float, start: float, stop: float) -> float:
    return max(value, stop) if (start > stop) else min(value, stop)


@registry.schedules("plateau.v1")
def plateau(
    max_patience: int, scale: float, schedule: Schedule[float]
) -> Schedule[float]:

    """Yields values from the wrapped schedule, exponentially scaled by the
    number of times optimization has plateaued. The caller must pass model
    evaluation scores through the last_score argument for the scaling to be
    adjusted. The last evaluation score is passed through the last_score argument
    as a tuple (last_score_step, last_score). This tuple indicates when a model
    was last evaluated (last_score_step) and with what score (last_score).

    max_patience (int): the number of evaluations without improvement when
        we consider the model to have plateaued.
    scale (float): scaling of the inner schedule (scale**n_plateaus * inner).
    schedule (Schedule[float]): the schedule to wrap.
    """

    return Schedule(
        "plateau",
        _plateau_schedule,
        attrs={
            "scale": scale,
            "max_patience": max_patience,
            "schedule": schedule,
            "state": _PlateauState(
                best_score=None, last_score_step=None, patience=0, n_plateaus=0
            ),
        },
    )


def _plateau_schedule(
    schedule: Schedule,
    step: int,
    *,
    last_score: Optional[Tuple[int, float]] = None,
    **kwargs,
) -> float:
    inner_schedule: Schedule[float] = schedule.attrs["schedule"]
    max_patience: int = schedule.attrs["max_patience"]
    scale: float = schedule.attrs["scale"]
    state: _PlateauState = schedule.attrs["state"]

    if last_score is None:
        return (scale**state.n_plateaus) * inner_schedule(
            step=step, last_score=last_score, **kwargs
        )

    last_score_step, last_score_ = last_score

    if (
        state.best_score is None
        or state.last_score_step is None
        or last_score_ > state.best_score
    ):
        state.best_score = last_score_
        state.patience = 0
    elif last_score_step < state.last_score_step:
        raise ValueError(
            f"Expected score with step >= {state.last_score_step}, was: {last_score_step}"
        )
    elif last_score_step > state.last_score_step:
        # If the score didn't improve and we are not seeing the last
        # score again, we may be at a plateau, so increase patience.
        state.patience += 1

        # If we are at the maximum patience, we consider the optimization
        # to have reached a plateau.
        if state.patience == max_patience:
            state.n_plateaus += 1
            state.patience = 0

    state.last_score_step = last_score_step

    return (scale**state.n_plateaus) * inner_schedule(
        step=step, last_score=last_score, **kwargs
    )


@dataclass
class _PlateauState:
    """Plateau schedule state.

    best_score (Optional[float]): the best score so far, or None when no
        score has been observed.
    last_score_step (Optional[int]): the step of the last score that was
        observed.
    patience (int): the number of scores so far which do not improve over
        the best score (reset after reaching the maximum patience).
    n_plateaus (int): the number of times the maximum patience has been
        reached.
    """

    best_score: Optional[float]
    last_score_step: Optional[int]
    patience: int
    n_plateaus: int

    # @dataclass(slots=True) is only supported in Python >= 3.10
    __slots__ = ["best_score", "last_score_step", "patience", "n_plateaus"]


@registry.schedules("slanted_triangular.v1")
def slanted_triangular(
    max_rate: float,
    num_steps: int,
    *,
    cut_frac: float = 0.1,
    ratio: int = 32,
    t: float = 0.0,
) -> Schedule[float]:
    """Yield an infinite series of values according to Howard and Ruder's
    "slanted triangular learning rate" schedule.
    """
    cut = int(num_steps * cut_frac)
    return Schedule(
        "slanted_triangular",
        _slanted_triangular_schedule,
        attrs={
            "max_rate": max_rate,
            "cut": cut,
            "cut_frac": cut_frac,
            "ratio": ratio,
            "t": t,
        },
    )


def _slanted_triangular_schedule(schedule: Schedule, step: int, **kwargs) -> float:
    max_rate = schedule.attrs["max_rate"]
    cut = schedule.attrs["cut"]
    cut_frac = schedule.attrs["cut_frac"]
    ratio = schedule.attrs["ratio"]
    t = schedule.attrs["t"]

    t_step = step + t + 1.0
    if t_step < cut:
        p = t_step / cut
    else:
        p = 1 - ((t_step - cut) / (cut * (1 / cut_frac - 1)))
    return max_rate * (1 + p * (ratio - 1)) * (1 / ratio)


@registry.schedules("warmup_linear.v1")
def warmup_linear(
    initial_rate: float, warmup_steps: int, total_steps: int
) -> Schedule[float]:
    """Generate a series, starting from an initial rate, and then with a warmup
    period, and then a linear decline. Used for learning rates.
    """
    return Schedule(
        "warmup_linear",
        _warmup_linear_schedule,
        attrs={
            "initial_rate": initial_rate,
            "warmup_steps": warmup_steps,
            "total_steps": total_steps,
        },
    )


def _warmup_linear_schedule(schedule: Schedule, step: int, **kwargs) -> float:
    initial_rate = schedule.attrs["initial_rate"]
    warmup_steps = schedule.attrs["warmup_steps"]
    total_steps = schedule.attrs["total_steps"]

    if step < warmup_steps:
        factor = step / max(1, warmup_steps)
    else:
        factor = max(0.0, (total_steps - step) / max(1.0, total_steps - warmup_steps))
    return factor * initial_rate


@registry.schedules("cyclic_triangular.v1")
def cyclic_triangular(min_lr: float, max_lr: float, period: int) -> Schedule[float]:
    return Schedule(
        "cyclic_triangular",
        _cyclic_triangular_schedule,
        attrs={"min_lr": min_lr, "max_lr": max_lr, "period": period},
    )


def _cyclic_triangular_schedule(schedule: Schedule, step: int, **kwargs) -> float:
    min_lr = schedule.attrs["min_lr"]
    max_lr = schedule.attrs["max_lr"]
    period = schedule.attrs["period"]

    it = step + 1
    # https://towardsdatascience.com/adaptive-and-cyclical-learning-rates-using-pytorch-2bf904d18dee
    cycle = numpy.floor(1 + it / (2 * period))
    x = numpy.abs(it / period - 2 * cycle + 1)
    relative = max(0, 1 - x)
    return min_lr + (max_lr - min_lr) * relative


__all__ = [
    "cyclic_triangular",
    "warmup_linear",
    "constant",
    "constant_then",
    "decaying",
    "warmup_linear",
    "slanted_triangular",
    "compounding",
]