"""
kiri/optim/optimizers.py
"""

import numpy as np
from kiri.backend.detect import BACKEND

if BACKEND == "mlx":
    import mlx.core as mx
    import mlx.optimizers as mlx_optim


# ─── CPU optimizers ───────────────────────────────────────────────────────────

class SGD:
    def __init__(self, params, lr=0.01, momentum=0.0, weight_decay=0.0):
        self.params = params
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self._velocities = [np.zeros_like(p.data) for p in params] if BACKEND != "mlx" else None
        self._mlx_opt = mlx_optim.SGD(learning_rate=lr, momentum=momentum) if BACKEND == "mlx" else None

    def zero_grad(self):
        if BACKEND != "mlx":
            for p in self.params:
                p.grad = np.zeros_like(p.data)

    def step(self, grads=None):
        if BACKEND == "mlx":
            # grads is a pytree of MLX arrays returned by mlx.grad()
            if grads is None:
                raise ValueError("MLX backend requires grads from mx.grad()")
            self._mlx_opt.update(self.params, grads)
            mx.eval(self.params)
        else:
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                g = p.grad
                if self.weight_decay != 0:
                    g = g + self.weight_decay * p.data
                if self.momentum != 0:
                    self._velocities[i] = self.momentum * self._velocities[i] + g
                    g = self._velocities[i]
                p.data -= self.lr * g


class Adam:
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999),
                 eps=1e-8, weight_decay=0.0):
        self.params = params
        self.lr = lr
        self.betas = betas
        self.eps = eps
        self.weight_decay = weight_decay
        self.t = 0
        self._mlx_opt = mlx_optim.Adam(learning_rate=lr, betas=betas, eps=eps) if BACKEND == "mlx" else None

        if BACKEND != "mlx":
            self._m = [np.zeros_like(p.data) for p in params]
            self._v = [np.zeros_like(p.data) for p in params]

    def zero_grad(self):
        if BACKEND != "mlx":
            for p in self.params:
                p.grad = np.zeros_like(p.data)

    def step(self, grads=None):
        if BACKEND == "mlx":
            if grads is None:
                raise ValueError("MLX backend requires grads from mx.grad()")
            self._mlx_opt.update(self.params, grads)
            mx.eval(self.params)
        else:
            self.t += 1
            b1, b2 = self.betas
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                g = p.grad
                if self.weight_decay != 0:
                    g = g + self.weight_decay * p.data
                self._m[i] = b1 * self._m[i] + (1 - b1) * g
                self._v[i] = b2 * self._v[i] + (1 - b2) * g**2
                m_hat = self._m[i] / (1 - b1**self.t)
                v_hat = self._v[i] / (1 - b2**self.t)
                p.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)


class AdamW(Adam):
    """Adam with decoupled weight decay (better for transformers)."""
    def step(self, grads=None):
        if BACKEND == "mlx":
            if grads is None:
                raise ValueError("MLX backend requires grads from mx.grad()")
            # Use MLX AdamW
            import mlx.optimizers as mlx_optim
            opt = mlx_optim.AdamW(learning_rate=self.lr,
                                   betas=self.betas,
                                   eps=self.eps,
                                   weight_decay=self.weight_decay)
            opt.update(self.params, grads)
            mx.eval(self.params)
        else:
            self.t += 1
            b1, b2 = self.betas
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                g = p.grad
                self._m[i] = b1 * self._m[i] + (1 - b1) * g
                self._v[i] = b2 * self._v[i] + (1 - b2) * g**2
                m_hat = self._m[i] / (1 - b1**self.t)
                v_hat = self._v[i] / (1 - b2**self.t)
                p.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
                p.data -= self.lr * self.weight_decay * p.data
