from typing import Callable
import numpy as np
from energizer.neural_network import Optimizer

class SGD(Optimizer):
    def __init__(self, params, lr: float = 0.01, momentum: float = 0.0, dampening: float = 0,  weight_decay: float = 0.0, nesterov: bool = False):
        default = {
            'lr': lr,
            'momentum': momentum,
            'weight_decay': weight_decay,
            'dampening': dampening,
            'nesterov': nesterov
        }
        super().__init__(params, default)

    def __setstate__(self, state):
        super().__setstate__(state)
        for group in self.param_groups:
            group.setdefault('nesterov', False)
            group.setdefault('dampening', 0)

    def step(self, closure: Callable = None):
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            lr = group['lr']
            momentum = group['momentum']
            dampening = group['dampening']
            nesterov = group['nesterov']
            weight_decay = group['weight_decay']

            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad.data if hasattr(p.grad, 'data') else p.grad
                
                p_data = p.data
                is_mlx = hasattr(p_data, '__class__') and 'mlx' in str(type(p_data))
                
                if is_mlx:
                    try:
                        import mlx.core as mx
                        if not (hasattr(grad, '__class__') and 'mlx' in str(type(grad))):
                            grad = mx.array(np.array(grad))
                    except ImportError:
                        grad = np.array(grad)
                else:
                    if not isinstance(grad, np.ndarray):
                        grad = np.array(grad)

                if weight_decay != 0:
                    grad = grad + weight_decay * p.data

                if momentum != 0:
                    if p not in self.state:
                        self.state[p] = {}
                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        if is_mlx:
                            import mlx.core as mx
                            buf = param_state['momentum_buffer'] = mx.zeros_like(p_data)
                        else:
                            buf = param_state['momentum_buffer'] = np.zeros_like(p_data)
                    else:
                        buf = param_state['momentum_buffer']
                    buf = momentum * buf + (1 - dampening) * grad
                    param_state['momentum_buffer'] = buf
                    if nesterov:
                        grad = grad + momentum * buf
                    else:
                        grad = buf
                p.data = p_data - lr * grad
        return loss

    
