import torch


class MissingClosureError(Exception):
    pass


class SAM(torch.optim.Optimizer):
    base_optimizer_cls = None

    def __init__(self, params, base_optimizer=None, rho=0.05, adaptive=False, **kwargs):
        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"

        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
        super(SAM, self).__init__(params, defaults)
        if base_optimizer is not None:
            self.base_optimizer_cls = base_optimizer
        assert self.base_optimizer_cls is not None, "Must provide base_optimizer_cls"
        self.base_optimizer = self.base_optimizer_cls(self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups

    @torch.no_grad()
    def first_step(self, zero_grad=False):
        grad_norm = self._grad_norm()
        for group in self.param_groups:
            scale = group["rho"] / (grad_norm + 1e-12)

            for p in group["params"]:
                if p.grad is None:
                    continue
                e_w = (
                    (torch.pow(p, 2) if group["adaptive"] else 1.0)
                    * p.grad
                    * scale.to(p)
                )
                p.add_(e_w)  # climb to the local maximum "w + e(w)"
                self.state[p]["e_w"] = e_w

        if zero_grad:
            self.zero_grad()

    @torch.no_grad()
    def second_step(self, zero_grad=False):
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                p.sub_(self.state[p]["e_w"])  # get back to "w" from "w + e(w)"

        self.base_optimizer.step()  # do the actual "sharpness-aware" update

        if zero_grad:
            self.zero_grad()

    @torch.no_grad()
    def step(self, closure=None):
        if closure is None:
            raise MissingClosureError(
                "Sharpness Aware Minimization requires closure, but it was not provided"
            )
        closure = torch.enable_grad()(
            closure
        )  # the closure should do a full forward-backward pass

        self.first_step(zero_grad=True)
        closure()
        self.second_step()

    def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][
            0
        ].device  # put everything on the same device, in case of model parallelism
        norm = torch.norm(
            torch.stack(
                [
                    ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad)
                    .norm(p=2)
                    .to(shared_device)
                    for group in self.param_groups
                    for p in group["params"]
                    if p.grad is not None
                ]
            ),
            p=2,  # type: ignore
        )
        return norm


class SAMOptimizer(SAM):
    """This is the class to be used in the pipeline.

    First you need to initialise this class
    with the torch optimiser
    that you would like to apply SAM to.

    Then, at the following call, the instance
    behaves extcly like the SAM optimiser

    Args:
        optimizer (torch.optim):
            the optimiser class (not the instance)
    """

    def __new__(cls, optimizer):
        return type(
            "SAMOptimizer",
            (SAM, torch.optim.Optimizer),
            {"base_optimizer_cls": optimizer},
        )
