#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

import torch
from pytext.config import ConfigBase
from pytext.config.component import Component, ComponentType


class Optimizer(Component):
    __COMPONENT_TYPE__ = ComponentType.OPTIMIZER
    __EXPANSIBLE__ = True

    class Config(ConfigBase):
        pass

    def backward(self, loss):
        loss.backward()

    def clip_grad_norm(self, model, max_norm):
        if max_norm is not None:
            return torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        else:
            return None

    def pre_export(self, model):
        pass

    def finalize(self) -> bool:
        return False


class Adagrad(torch.optim.Adagrad, Optimizer):
    class Config(Optimizer.Config):
        lr: float = 1e-2
        weight_decay: float = 0.00001

    def __init__(self, parameters, lr, weight_decay):
        super().__init__(parameters, lr=lr, weight_decay=weight_decay)

    @classmethod
    def from_config(cls, config: Config, model: torch.nn.Module):
        return cls(model.parameters(), config.lr, config.weight_decay)


class Adam(torch.optim.Adam, Optimizer):
    class Config(Optimizer.Config):
        lr: float = 0.001
        weight_decay: float = 0.00001
        eps: float = 1e-8

    def __init__(self, parameters, lr, weight_decay, eps):
        super().__init__(parameters, lr=lr, weight_decay=weight_decay, eps=eps)

    @classmethod
    def from_config(cls, config: Config, model: torch.nn.Module):
        return cls(model.parameters(), config.lr, config.weight_decay, config.eps)


class SGD(torch.optim.SGD, Optimizer):
    class Config(Optimizer.Config):
        lr: float = 0.001
        momentum: float = 0.0

    def __init__(self, parameters, lr, momentum):
        super().__init__(parameters, lr=lr, momentum=momentum)

    @classmethod
    def from_config(cls, config: Config, model: torch.nn.Module):
        return cls(model.parameters(), config.lr, config.momentum)


class AdamW(torch.optim.AdamW, Optimizer):
    """Adds PyText support for
       Decoupled Weight Decay Regularization for Adam as done in the paper:
       https://arxiv.org/abs/1711.05101
       for more information read the fast.ai blog on this optimization
       method here: https://www.fast.ai/2018/07/02/adam-weight-decay/
    """

    class Config(Optimizer.Config):
        lr: float = 0.001
        weight_decay: float = 1e-2
        eps: float = 1e-8

    def __init__(self, parameters, lr, weight_decay, eps):
        super().__init__(parameters, lr=lr, weight_decay=weight_decay, eps=eps)

    @classmethod
    def from_config(cls, config: Config, model: torch.nn.Module):
        return cls(model.parameters(), config.lr, config.weight_decay, config.eps)


def learning_rates(optimizer):
    for param_group in optimizer.param_groups:
        yield param_group["lr"]
