import torch
import torch.nn as nn
from torch.optim.optimizer import Optimizer
from typing import Any, Dict, List, Optional, Tuple, Callable
import math
from collections import defaultdict
import random


class SAM(Optimizer):
    """Sharpness-Aware Minimization optimizer"""

    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}"
        
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD

        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups

    @torch.no_grad()
    def first_step(self, zero_grad=False):
        """First step: compute adversarial parameters"""
        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
                self.state[p]["old_p"] = p.data.clone()
                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)"

        if zero_grad:
            self.zero_grad()

    @torch.no_grad()
    def second_step(self, zero_grad=False):
        """Second step: update parameters using base optimizer"""
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                p.data = self.state[p]["old_p"]  # 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):
        """Combined step function"""
        assert closure is not None, "SAM requires a closure, but it was not provided"

        # first forward-backward pass
        self.first_step(zero_grad=True)

        # second forward-backward pass
        with torch.enable_grad():
            closure()
        self.second_step()

    def _grad_norm(self):
        """Compute the norm of gradients"""
        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(dtype=torch.float32).to(shared_device)
                    for group in self.param_groups
                    for p in group["params"]
                    if p.grad is not None
                ]
            ),
            dtype=torch.float32,
        )
        return norm


class GSAM(Optimizer):
    """Gradient-based Sharpness-Aware Minimization"""

    def __init__(self, params, base_optimizer=None, rho=0.05, alpha=0.4, adaptive=False, **kwargs):
        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
        assert alpha >= 0.0, f"Invalid alpha, should be non-negative: {alpha}"
        
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD

        defaults = dict(rho=rho, alpha=alpha, adaptive=adaptive, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(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
                self.state[p]["old_p"] = p.data.clone()
                self.state[p]["old_g"] = p.grad.clone()

                e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale.to(p)
                p.add_(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

                # gradient decomposition
                old_g = self.state[p]["old_g"]
                new_g = p.grad

                inner_prod = torch.sum(old_g * new_g)
                old_g_norm_sq = torch.sum(old_g * old_g)

                # cosine similarity based decomposition
                if old_g_norm_sq > 0:
                    cos_sim = inner_prod / old_g_norm_sq
                    cos_sim = torch.clamp(cos_sim, 0, 1)  # ensure non-negative

                    # surrogate gap and gradient
                    surrogate_g = cos_sim * old_g + group["alpha"] * (new_g - cos_sim * old_g)
                    p.grad = surrogate_g

                p.data = self.state[p]["old_p"]

        self.base_optimizer.step()

        if zero_grad:
            self.zero_grad()

    def step(self, closure=None):
        assert closure is not None, "GSAM requires a closure"

        self.first_step(zero_grad=True)

        with torch.enable_grad():
            closure()
        self.second_step()

    def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][0].device
        norm = torch.norm(
            torch.stack(
                [
                    ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(dtype=torch.float32).to(shared_device)
                    for group in self.param_groups
                    for p in group["params"]
                    if p.grad is not None
                ]
            ),
            dtype=torch.float32,
        )
        return norm


class ASAM(Optimizer):
    """Adaptive Sharpness-Aware Minimization"""

    def __init__(self, params, base_optimizer=None, rho=0.5, eta=0.01, adaptive=True, **kwargs):
        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
        
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD

        defaults = dict(rho=rho, eta=eta, adaptive=adaptive, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(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

                self.state[p]["old_p"] = p.data.clone()

                if group["adaptive"]:
                    # adaptive scaling based on parameter magnitude
                    abs_p = torch.abs(p.data)
                    k = torch.clamp(abs_p, min=group["eta"])
                    e_w = p.grad * scale.to(p) * k
                else:
                    e_w = p.grad * scale.to(p)

                p.add_(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.data = self.state[p]["old_p"]

        self.base_optimizer.step()

        if zero_grad:
            self.zero_grad()

    def step(self, closure=None):
        assert closure is not None, "ASAM requires a closure"

        self.first_step(zero_grad=True)

        with torch.enable_grad():
            closure()
        self.second_step()

    def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][0].device
        norm = torch.norm(
            torch.stack(
                [
                    p.grad.norm(dtype=torch.float32).to(shared_device)
                    for group in self.param_groups
                    for p in group["params"]
                    if p.grad is not None
                ]
            ),
            dtype=torch.float32,
        )
        return norm


class LookSAM(Optimizer):
    """Look-ahead SAM optimizer"""

    def __init__(self, params, base_optimizer=None, k=5, alpha=0.5, rho=0.05, **kwargs):
        assert 0 <= alpha <= 1, f"Invalid alpha, should be in [0, 1]: {alpha}"
        assert k >= 1, f"Invalid k, should be >= 1: {k}"
        
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD

        defaults = dict(k=k, alpha=alpha, rho=rho, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups

        self.step_count = 0

    @torch.no_grad()
    def lookahead_step(self):
        """Perform lookahead update"""
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue

                param_state = self.state[p]
                if "slow_buffer" not in param_state:
                    param_state["slow_buffer"] = torch.zeros_like(p.data)
                    param_state["slow_buffer"].copy_(p.data)

                slow = param_state["slow_buffer"]
                slow.add_(p.data - slow, alpha=group["alpha"])
                p.data.copy_(slow)

    def step(self, closure=None):
        assert closure is not None, "LookSAM requires a closure"

        # SAM step
        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
                if "old_p" not in self.state[p]:
                    self.state[p]["old_p"] = torch.zeros_like(p.data)
                self.state[p]["old_p"].copy_(p.data)
                p.add_(p.grad * scale.to(p))

        self.zero_grad()

        with torch.enable_grad():
            closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                p.data.copy_(self.state[p]["old_p"])

        self.base_optimizer.step()

        self.step_count += 1

        # Lookahead step
        if self.step_count % self.param_groups[0]["k"] == 0:
            self.lookahead_step()

        self.zero_grad()

    def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][0].device
        norm = torch.norm(
            torch.stack(
                [
                    p.grad.norm(dtype=torch.float32).to(shared_device)
                    for group in self.param_groups
                    for p in group["params"]
                    if p.grad is not None
                ]
            ),
            dtype=torch.float32,
        )
        return norm


class WSAM(Optimizer):
    """Weighted Sharpness-Aware Minimization"""

    def __init__(self, params, base_optimizer=None, rho=0.05, tau=1.0, **kwargs):
        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
        
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD

        defaults = dict(rho=rho, tau=tau, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(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

                self.state[p]["old_p"] = p.data.clone()

                # compute importance weights
                grad_norm_p = p.grad.norm()
                weight = torch.exp(group["tau"] * grad_norm_p)

                e_w = weight * p.grad * scale.to(p)
                p.add_(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.data = self.state[p]["old_p"]

        self.base_optimizer.step()

        if zero_grad:
            self.zero_grad()

    def step(self, closure=None):
        assert closure is not None, "WSAM requires a closure"

        self.first_step(zero_grad=True)

        with torch.enable_grad():
            closure()
        self.second_step()

    def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][0].device
        norm = torch.norm(
            torch.stack(
                [
                    p.grad.norm(dtype=torch.float32).to(shared_device)
                    for group in self.param_groups
                    for p in group["params"]
                    if p.grad is not None
                ]
            ),
            dtype=torch.float32,
        )
        return norm


class GradientCentralization(Optimizer):
    """Gradient Centralization wrapper"""

    def __init__(self, params, base_optimizer=None, use_gc=True, gc_conv_only=False, **kwargs):
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD
            
        defaults = dict(use_gc=use_gc, gc_conv_only=gc_conv_only, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups

    def centralize_gradient(self, p):
        """Apply gradient centralization"""
        if len(p.shape) > 1:  # only for conv and linear layers
            p.grad.data = p.grad.data - p.grad.data.mean(dim=tuple(range(1, len(p.shape))), keepdim=True)

    def step(self, closure=None):
        # Apply gradient centralization
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue

                if group["use_gc"]:
                    if group["gc_conv_only"]:
                        # only apply to conv layers (4D tensors)
                        if len(p.shape) == 4:
                            self.centralize_gradient(p)
                    else:
                        # apply to all multi-dimensional parameters
                        if len(p.shape) > 1:
                            self.centralize_gradient(p)

        # Call base optimizer
        return self.base_optimizer.step(closure)


class PCGrad(Optimizer):
    """Projecting Conflicting Gradients"""

    def __init__(self, params, base_optimizer=None, num_tasks=1, **kwargs):
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD
            
        defaults = dict(num_tasks=num_tasks, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups
        self.num_tasks = num_tasks

    def project_conflicting_gradients(self, grads):
        """Project conflicting gradients"""
        pc_grads = []

        for i in range(len(grads)):
            grad_i = grads[i]
            pc_grad = grad_i.clone()

            for j in range(len(grads)):
                if i != j:
                    grad_j = grads[j]

                    # Ensure gradients have the same shape for inner product
                    if grad_i.shape != grad_j.shape:
                        # Use the minimum size for projection
                        min_size = min(grad_i.shape[0], grad_j.shape[0])
                        grad_i_proj = grad_i[:min_size]
                        grad_j_proj = grad_j[:min_size]
                    else:
                        grad_i_proj = grad_i
                        grad_j_proj = grad_j

                    # compute inner product
                    inner_prod = torch.sum(grad_i_proj * grad_j_proj)

                    # if negative inner product (conflicting), project
                    if inner_prod < 0:
                        grad_j_norm_sq = torch.sum(grad_j_proj * grad_j_proj)
                        if grad_j_norm_sq > 0:
                            # Apply projection to the truncated version
                            projection = (inner_prod / grad_j_norm_sq) * grad_j_proj
                            
                            # Pad or truncate projection to match original grad_i shape
                            if grad_i.shape != grad_j.shape:
                                if grad_i.shape[0] > grad_j.shape[0]:
                                    # Pad with zeros
                                    padding = torch.zeros(grad_i.shape[0] - min_size, device=grad_i.device, dtype=grad_i.dtype)
                                    projection = torch.cat([projection, padding])
                                else:
                                    # Truncate
                                    projection = projection[:grad_i.shape[0]]
                            
                            pc_grad = pc_grad - projection

            pc_grads.append(pc_grad)

        return pc_grads

    def step(self, closure=None, task_losses=None):
        """
        Args:
            closure: closure function
            task_losses: list of task-specific loss tensors or functions
        """
        if task_losses is not None:
            # compute gradients for each task
            task_grads = []
            for task_loss in task_losses:
                self.zero_grad()
                
                # Handle both loss tensors and loss functions
                if callable(task_loss):
                    loss_tensor = task_loss()
                else:
                    loss_tensor = task_loss
                    
                loss_tensor.backward(retain_graph=True)

                # collect gradients
                grads = []
                for group in self.param_groups:
                    for p in group["params"]:
                        if p.grad is not None:
                            grads.append(p.grad.clone().flatten())

                if grads:
                    task_grads.append(torch.cat(grads))

            # project conflicting gradients
            if task_grads:
                pc_grads = self.project_conflicting_gradients(task_grads)
                
                # Ensure all gradients have the same size before averaging
                if len(pc_grads) > 1:
                    # Find the minimum size
                    min_size = min(grad.shape[0] for grad in pc_grads)
                    
                    # Truncate all gradients to the same size
                    truncated_grads = [grad[:min_size] for grad in pc_grads]
                    avg_grad = sum(truncated_grads) / len(truncated_grads)
                    
                    # Pad back to the original size of the first gradient
                    if pc_grads[0].shape[0] > min_size:
                        padding = torch.zeros(pc_grads[0].shape[0] - min_size, device=avg_grad.device, dtype=avg_grad.dtype)
                        avg_grad = torch.cat([avg_grad, padding])
                else:
                    avg_grad = pc_grads[0]

                # assign back to parameters
                idx = 0
                for group in self.param_groups:
                    for p in group["params"]:
                        if p.grad is not None:
                            param_size = p.numel()
                            p.grad.data = avg_grad[idx : idx + param_size].reshape(p.shape)
                            idx += param_size

        return self.base_optimizer.step(closure)


class GradNorm(Optimizer):
    """Gradient Normalization for multi-task learning"""

    def __init__(self, params, base_optimizer=None, num_tasks=1, alpha=1.5, **kwargs):
        # Default to SGD if no base_optimizer provided
        if base_optimizer is None:
            import torch.optim as optim
            base_optimizer = optim.SGD
            
        defaults = dict(num_tasks=num_tasks, alpha=alpha, **kwargs)
        super().__init__(params, defaults)

        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups
        self.num_tasks = num_tasks
        self.task_weights = nn.Parameter(torch.ones(num_tasks))

    def step(self, closure=None, task_losses=None, task_grads=None):
        """
        Args:
            closure: closure function
            task_losses: tensor of task losses
            task_grads: list of task gradient norms
        """
        if task_losses is not None and task_grads is not None:
            # compute relative task losses
            loss_ratios = task_losses / task_losses[0]  # relative to first task

            # compute gradient norm ratios
            grad_norm_avg = sum(task_grads) / len(task_grads)
            grad_ratios = [g / grad_norm_avg for g in task_grads]

            # compute targets
            targets = []
            for i in range(self.num_tasks):
                if i == 0:
                    targets.append(torch.tensor(1.0))
                else:
                    target = loss_ratios[i] ** self.param_groups[0]["alpha"]
                    targets.append(target)

            # update task weights
            for i in range(self.num_tasks):
                target = targets[i]
                grad_ratio = grad_ratios[i]

                # gradnorm loss
                gradnorm_loss = torch.abs(grad_ratio - target.detach())

                # update weight
                weight_grad = torch.autograd.grad(gradnorm_loss, self.task_weights[i], retain_graph=True)[0]
                self.task_weights.data[i] -= 0.025 * weight_grad  # small lr for weights

            # renormalize weights
            self.task_weights.data = self.num_tasks * self.task_weights.data / self.task_weights.data.sum()

        return self.base_optimizer.step(closure)


__all__ = ["SAM", "GSAM", "ASAM", "LookSAM", "WSAM", "GradientCentralization", "PCGrad", "GradNorm"]
