import mx

class SGD:
    def __init__(self, parameters, learning_rate=0.01):
        self.params = parameters
        self.lr = learning_rate
    
    def step(self):
        """Update parameters using their gradients"""
        for param in self.params:
            if param.grad is not None:
                # param = param - lr * grad (in-place would be better)
                # For now, create new tensor
                updated = param - param.grad.scalar_mul(self.lr)
                # Copy data back (this is a workaround)
                param.data = updated.data
    
    def zero_grad(self):
        """Zero out all gradients"""
        for param in self.params:
            param.zero_grad()
    
class Adam:
    def __init__(self, parameters, learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.params = parameters
        self.lr = learning_rate
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.m = {}  # First moment
        self.v = {}  # Second moment
        self.t = 0   # Timestep
    
    def step(self):
        """Update parameters using their gradients"""
        self.t += 1
        
        for i, param in enumerate(self.params):
            if param.grad is None:
                continue
            
            # Initialize moments if first time
            if i not in self.m:
                self.m[i] = mx.zeros(param.shape)
                self.v[i] = mx.zeros(param.shape)
            
            # Update biased moments
            self.m[i] = self.m[i].scalar_mul(self.beta1) + param.grad.scalar_mul(1 - self.beta1)
            self.v[i] = self.v[i].scalar_mul(self.beta2) + (param.grad * param.grad).scalar_mul(1 - self.beta2)
            
            # Bias correction
            m_hat = self.m[i].scalar_div(1 - self.beta1 ** self.t)
            v_hat = self.v[i].scalar_div(1 - self.beta2 ** self.t)
            
            # Update parameters (simplified - would need sqrt and div operations)
            # param.data = param.data - self.lr * m_hat / (sqrt(v_hat) + self.eps)
            # For now, just do SGD-style update
            updated = param - m_hat.scalar_mul(self.lr)
            param.data = updated.data
    
    def zero_grad(self):
        """Zero out all gradients"""
        for param in self.params:
            param.zero_grad()