"""
Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""

from __future__ import annotations

import inspect
import math
from bisect import bisect

import torch

from fairchem.core.common.typing import assert_is_instance as aii


def multiply(obj, num):
    if isinstance(obj, list):
        for i in range(len(obj)):
            obj[i] = obj[i] * num
    else:
        obj = obj * num
    return obj


def cosine_lr_lambda(current_step: int, scheduler_params):
    warmup_epochs = aii(scheduler_params["warmup_epochs"], int)
    lr_warmup_factor = aii(scheduler_params["warmup_factor"], float)
    max_epochs = aii(scheduler_params["epochs"], int)
    lr_min_factor = aii(scheduler_params["lr_min_factor"], float)

    # `warmup_epochs` is already multiplied with the num of iterations
    if current_step <= warmup_epochs:
        alpha = current_step / float(warmup_epochs)
        return lr_warmup_factor * (1.0 - alpha) + alpha
    else:
        if current_step >= max_epochs:
            return lr_min_factor
        return lr_min_factor + 0.5 * (1 - lr_min_factor) * (
            1 + math.cos(math.pi * (current_step / max_epochs))
        )


class CosineLRLambda:
    def __init__(self, scheduler_params) -> None:
        self.warmup_epochs = aii(scheduler_params["warmup_epochs"], int)
        self.lr_warmup_factor = aii(scheduler_params["warmup_factor"], float)
        self.max_epochs = aii(scheduler_params["epochs"], int)
        self.lr_min_factor = aii(scheduler_params["lr_min_factor"], float)

    def __call__(self, current_step: int):
        # `warmup_epochs` is already multiplied with the num of iterations
        if current_step <= self.warmup_epochs:
            alpha = current_step / float(self.warmup_epochs)
            return self.lr_warmup_factor * (1.0 - alpha) + alpha
        else:
            if current_step >= self.max_epochs:
                return self.lr_min_factor
            return self.lr_min_factor + 0.5 * (1 - self.lr_min_factor) * (
                1 + math.cos(math.pi * (current_step / self.max_epochs))
            )


def multistep_lr_lambda(current_step: int, scheduler_params) -> float:
    warmup_epochs = aii(scheduler_params["warmup_epochs"], int)
    lr_warmup_factor = aii(scheduler_params["warmup_factor"], float)
    lr_decay_epochs: list[int] = scheduler_params["decay_epochs"]
    lr_gamma = aii(scheduler_params["decay_rate"], float)

    if current_step <= warmup_epochs:
        alpha = current_step / float(warmup_epochs)
        return lr_warmup_factor * (1.0 - alpha) + alpha
    else:
        idx = bisect(lr_decay_epochs, current_step)
        return pow(lr_gamma, idx)


class MultistepLRLambda:
    def __init__(self, scheduler_params) -> None:
        self.warmup_epochs = aii(scheduler_params["warmup_epochs"], int)
        self.lr_warmup_factor = aii(scheduler_params["warmup_factor"], float)
        self.lr_decay_epochs = aii(scheduler_params["decay_epochs"], list)
        self.lr_gamma = aii(scheduler_params["decay_rate"], float)

    def __call__(self, current_step: int) -> float:
        if current_step <= self.warmup_epochs:
            alpha = current_step / float(self.warmup_epochs)
            return self.lr_warmup_factor * (1.0 - alpha) + alpha
        else:
            idx = bisect(self.lr_decay_epochs, current_step)
            return pow(self.lr_gamma, idx)


class LRScheduler:
    """
    Notes:
        1. scheduler.step() is called for every step for OC20 training.
        2. We use "scheduler_params" in .yml to specify scheduler parameters.
        3. For cosine learning rate, we use LambdaLR with lambda function being cosine:
            scheduler: LambdaLR
            scheduler_params:
                lambda_type: cosine
                ...
        4. Following 3., if `cosine` is used, `scheduler_params` in .yml looks like:
            scheduler: LambdaLR
            scheduler_params:
                lambda_type: cosine
                warmup_epochs: ...
                warmup_factor: ...
                lr_min_factor: ...
        5. Following 3., if `multistep` is used, `scheduler_params` in .yml looks like:
            scheduler: LambdaLR
            scheduler_params:
                lambda_type: multistep
                warmup_epochs: ...
                warmup_factor: ...
                decay_epochs: ... (list)
                decay_rate: ...

    Args:
        optimizer (obj): torch optim object
        config (dict): Optim dict from the input config
    """

    def __init__(self, optimizer, config) -> None:
        self.optimizer = optimizer
        self.config = config.copy()

        assert "scheduler" in self.config
        assert "scheduler_params" in self.config
        self.scheduler_type = aii(self.config["scheduler"], str)
        self.scheduler_params = self.config["scheduler_params"].copy()

        # Use `LambdaLR` for multi-step and cosine learning rate
        if self.scheduler_type == "LambdaLR":
            scheduler_lambda_fn = None
            self.lambda_type = self.scheduler_params["lambda_type"]

            if self.lambda_type == "cosine":
                scheduler_lambda_fn = CosineLRLambda(self.scheduler_params)
            elif self.lambda_type == "multistep":
                scheduler_lambda_fn = MultistepLRLambda(self.scheduler_params)
            else:
                raise ValueError
            self.scheduler_params["lr_lambda"] = scheduler_lambda_fn

        if self.scheduler_type != "Null":
            self.scheduler = getattr(torch.optim.lr_scheduler, self.scheduler_type)
            scheduler_args = self.filter_kwargs(self.scheduler_params)
            self.scheduler = self.scheduler(optimizer, **scheduler_args)

    def step(self, metrics=None, epoch=None):
        if self.scheduler_type == "Null":
            return
        if self.scheduler_type == "ReduceLROnPlateau":
            if metrics is None:
                raise Exception("Validation set required for ReduceLROnPlateau.")
            self.scheduler.step(metrics)
        else:
            self.scheduler.step()

    def filter_kwargs(self, config):
        # adapted from https://stackoverflow.com/questions/26515595/
        sig = inspect.signature(self.scheduler)
        filter_keys = [
            param.name
            for param in sig.parameters.values()
            if param.kind == param.POSITIONAL_OR_KEYWORD
        ]
        filter_keys.remove("optimizer")
        return {arg: config[arg] for arg in config if arg in filter_keys}

    def get_lr(self) -> float | None:
        for group in self.optimizer.param_groups:
            return aii(group["lr"], float)
        return None
