"""
Natural Gradient Descent optimizer using the Fisher Information Matrix.

Update rule:  theta_new = theta - eta * F^{-1} @ grad_L

The Fisher metric accounts for the information geometry of the parameter
space, yielding faster convergence than ordinary gradient descent (and
often faster than Adam) especially for fine-tuning and few-shot regimes.

Author: Sheng-Kai Huang <akai@fawstudio.com>
"""

from __future__ import annotations

from typing import Callable, Optional

import mlx.core as mx


class NaturalGradientOptimizer:
    """
    Natural gradient descent with online Fisher matrix estimation.

    Parameters
    ----------
    lr : float
        Learning rate eta.
    damping : float
        Tikhonov regularisation lambda added to F before inversion,
        i.e. we invert (F + lambda I). Prevents blow-up when F is
        near-singular.
    fisher_update_interval : int
        Re-estimate the Fisher matrix every this many steps.
        Set to 1 for exact natural gradient (expensive).
    ema_decay : float
        Exponential moving average decay for the Fisher matrix.
        F_t = decay * F_{t-1} + (1 - decay) * F_new.
        Set to 0 for no EMA (just use latest estimate).
    """

    def __init__(
        self,
        lr: float = 1e-3,
        damping: float = 1e-4,
        fisher_update_interval: int = 1,
        ema_decay: float = 0.0,
    ):
        self.lr = lr
        self.damping = damping
        self.fisher_update_interval = fisher_update_interval
        self.ema_decay = ema_decay

        self._F: Optional[mx.array] = None
        self._F_inv: Optional[mx.array] = None
        self._step_count: int = 0

    @property
    def fisher_matrix(self) -> Optional[mx.array]:
        """Current Fisher matrix estimate (or None if not yet computed)."""
        return self._F

    def step(
        self,
        theta: mx.array,
        grad: mx.array,
        fisher_estimator: Optional[Callable[[mx.array], mx.array]] = None,
    ) -> mx.array:
        """
        Perform one natural gradient step.

        Parameters
        ----------
        theta : mx.array, shape (n_params,)
            Current parameters.
        grad : mx.array, shape (n_params,)
            Ordinary (Euclidean) gradient of the loss w.r.t. theta.
        fisher_estimator : callable or None
            A function theta -> mx.array (n_params, n_params) that
            estimates the Fisher matrix at the given theta.
            Required on the first call and every ``fisher_update_interval``
            steps thereafter.

        Returns
        -------
        mx.array, shape (n_params,)
            Updated parameters theta_new.
        """
        n = theta.shape[0]

        # Update Fisher estimate if needed
        if (
            self._F is None
            or self._step_count % self.fisher_update_interval == 0
        ):
            if fisher_estimator is None:
                if self._F is None:
                    raise ValueError(
                        "fisher_estimator must be provided on the first call "
                        "(or whenever the cached Fisher is stale)."
                    )
                # Reuse cached F
            else:
                F_new = fisher_estimator(theta)
                if self._F is None or self.ema_decay == 0.0:
                    self._F = F_new
                else:
                    self._F = self.ema_decay * self._F + (1.0 - self.ema_decay) * F_new
                # Recompute inverse
                F_reg = self._F + self.damping * mx.eye(n, dtype=self._F.dtype)
                self._F_inv = mx.linalg.inv(F_reg, stream=mx.cpu)

        # Natural gradient step
        natural_grad = self._F_inv @ grad
        theta_new = theta - self.lr * natural_grad

        self._step_count += 1
        mx.eval(theta_new)
        return theta_new

    def reset(self):
        """Reset internal state (Fisher estimate and step counter)."""
        self._F = None
        self._F_inv = None
        self._step_count = 0


def natural_gradient_descent(
    loss_fn: Callable[[mx.array], mx.array],
    fisher_fn: Callable[[mx.array], mx.array],
    theta_init: mx.array,
    lr: float = 1e-2,
    damping: float = 1e-4,
    n_steps: int = 100,
    verbose: bool = False,
) -> tuple[mx.array, list[float]]:
    """
    Run natural gradient descent to minimize a loss function.

    Parameters
    ----------
    loss_fn : callable  theta -> scalar
        Loss function to minimize.
    fisher_fn : callable  theta -> mx.array (n, n)
        Fisher matrix estimator at theta.
    theta_init : mx.array, shape (n_params,)
        Initial parameters.
    lr : float
        Learning rate.
    damping : float
        Tikhonov regularisation for Fisher inversion.
    n_steps : int
        Number of optimisation steps.
    verbose : bool
        Print loss every 10 steps.

    Returns
    -------
    theta_final : mx.array
    loss_history : list of float
    """
    theta = mx.array(theta_init, dtype=mx.float32)
    n = theta.shape[0]
    loss_history = []

    for step in range(n_steps):
        # Compute loss and gradient via central differences
        loss_val = loss_fn(theta)
        mx.eval(loss_val)
        loss_history.append(loss_val.item())

        # Numerical gradient
        eps = 1e-4
        grad = mx.zeros_like(theta)
        grad_list = []
        for i in range(n):
            e_i = mx.zeros_like(theta)
            e_i[i] = eps
            g_i = (loss_fn(theta + e_i) - loss_fn(theta - e_i)) / (2.0 * eps)
            grad_list.append(g_i)
        grad = mx.stack(grad_list)

        # Fisher and natural gradient
        F = fisher_fn(theta)
        F_reg = F + damping * mx.eye(n, dtype=F.dtype)
        F_inv = mx.linalg.inv(F_reg, stream=mx.cpu)
        nat_grad = F_inv @ grad

        theta = theta - lr * nat_grad
        mx.eval(theta)

        if verbose and step % 10 == 0:
            print(f"  step {step:4d}  loss = {loss_val.item():.6e}")

    return theta, loss_history
