"""
Core Fisher Information Matrix computation using MLX.

Provides three methods:
  1. From a log-likelihood function (numerical differentiation on GPU)
  2. From model predictions + errors (analytic Gaussian Fisher)
  3. From samples (empirical Fisher via outer product of score vectors)

Plus specialized routines for CMB C_l Fisher matrices, KL divergence,
and quantum relative entropy.

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

from __future__ import annotations

from typing import Callable, Optional, Sequence

import mlx.core as mx
import numpy as np


class FisherMatrix:
    """
    Compute the Fisher information matrix:

        F_ij = E[ (d log p / d theta_i) * (d log p / d theta_j) ]

    Three construction methods are provided as class methods.
    """

    def __init__(self, matrix: mx.array):
        """
        Wrap a pre-computed Fisher matrix.

        Parameters
        ----------
        matrix : mx.array, shape (n_params, n_params)
        """
        if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
            raise ValueError("Fisher matrix must be square.")
        self._matrix = matrix

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def matrix(self) -> mx.array:
        """Return the Fisher matrix as an MLX array."""
        return self._matrix

    @property
    def n_params(self) -> int:
        return self._matrix.shape[0]

    def to_numpy(self) -> np.ndarray:
        """Convert to a NumPy array."""
        return np.array(self._matrix)

    # ------------------------------------------------------------------
    # Derived quantities
    # ------------------------------------------------------------------

    def inverse(self, reg: float = 0.0) -> mx.array:
        """
        Invert the Fisher matrix (= covariance matrix of the MLE).

        Parameters
        ----------
        reg : float
            Tikhonov regularisation added to the diagonal before inversion.
        """
        F = self._matrix
        if reg > 0:
            F = F + reg * mx.eye(self.n_params, dtype=F.dtype)
        return mx.linalg.inv(F, stream=mx.cpu)

    def marginal_errors(self, reg: float = 0.0) -> mx.array:
        """
        1-sigma marginal errors: sqrt(diag(F^{-1})).
        """
        cov = self.inverse(reg=reg)
        return mx.sqrt(mx.diag(cov))

    def eigenvalues(self) -> mx.array:
        """Eigenvalues of the Fisher matrix (sorted ascending)."""
        eigvals, _ = mx.linalg.eigh(self._matrix, stream=mx.cpu)
        return eigvals

    def condition_number(self) -> float:
        """Condition number = max(|eigenvalue|) / min(|eigenvalue|)."""
        ev = self.eigenvalues()
        ev_abs = mx.abs(ev)
        return (mx.max(ev_abs) / mx.max(mx.array([mx.min(ev_abs), mx.array(1e-30)]))).item()

    # ------------------------------------------------------------------
    # Construction: from log-likelihood
    # ------------------------------------------------------------------

    @classmethod
    def from_loglikelihood(
        cls,
        log_likelihood: Callable[[mx.array], mx.array],
        theta: mx.array,
        n_samples: int = 1000,
        eps: float = 1e-4,
    ) -> "FisherMatrix":
        """
        Estimate the Fisher matrix by numerical differentiation of a
        log-likelihood function, evaluated over *n_samples* data
        realisations drawn internally by the callable.

        The callable ``log_likelihood(theta)`` must return a 1-D array of
        shape ``(n_samples,)`` giving the log-likelihood for each sample.

        Parameters
        ----------
        log_likelihood : callable  theta -> mx.array (n_samples,)
        theta : mx.array, shape (n_params,)
        n_samples : int
            Expected first dimension of the output.
        eps : float
            Step size for central-difference gradient.

        Returns
        -------
        FisherMatrix
        """
        theta = mx.array(theta, dtype=mx.float32)
        n_params = theta.shape[0]

        # Compute score vectors via central differences: shape (n_params, n_samples)
        scores = []
        for i in range(n_params):
            e_i = mx.zeros_like(theta)
            e_i[i] = eps
            ll_plus = log_likelihood(theta + e_i)
            ll_minus = log_likelihood(theta - e_i)
            score_i = (ll_plus - ll_minus) / (2.0 * eps)
            scores.append(score_i)

        # Stack: (n_params, n_samples)
        S = mx.stack(scores, axis=0)

        # Fisher = (1/n) S @ S^T  (empirical expectation of outer product)
        n = S.shape[1]
        F = (S @ S.T) / n
        mx.eval(F)
        return cls(F)

    # ------------------------------------------------------------------
    # Construction: from model predictions + errors (Gaussian Fisher)
    # ------------------------------------------------------------------

    @classmethod
    def from_model(
        cls,
        model_fn: Callable[[mx.array], mx.array],
        theta: mx.array,
        sigma: mx.array,
        eps: float = 1e-4,
    ) -> "FisherMatrix":
        """
        Gaussian Fisher matrix:

            F_ij = sum_k  (d mu_k / d theta_i)(d mu_k / d theta_j) / sigma_k^2

        Parameters
        ----------
        model_fn : callable  theta -> mx.array (n_data,)
            Returns model predictions mu(theta).
        theta : mx.array, shape (n_params,)
            Fiducial parameter values.
        sigma : mx.array, shape (n_data,)
            1-sigma measurement errors.
        eps : float
            Step size for numerical derivatives.

        Returns
        -------
        FisherMatrix
        """
        theta = mx.array(theta, dtype=mx.float32)
        sigma = mx.array(sigma, dtype=mx.float32)
        n_params = theta.shape[0]

        # Jacobian d mu / d theta  via central differences: (n_data, n_params)
        jacobian_cols = []
        for i in range(n_params):
            e_i = mx.zeros_like(theta)
            e_i[i] = eps
            mu_plus = model_fn(theta + e_i)
            mu_minus = model_fn(theta - e_i)
            jacobian_cols.append((mu_plus - mu_minus) / (2.0 * eps))

        # J shape: (n_data, n_params)
        J = mx.stack(jacobian_cols, axis=-1)

        # Weight by 1/sigma^2
        # F = J^T @ diag(1/sigma^2) @ J
        inv_var = 1.0 / (sigma * sigma)  # (n_data,)
        J_weighted = J * inv_var[:, None]  # broadcast: (n_data, n_params)
        F = J_weighted.T @ J  # (n_params, n_params)
        mx.eval(F)
        return cls(F)

    # ------------------------------------------------------------------
    # Construction: from samples (empirical Fisher)
    # ------------------------------------------------------------------

    @classmethod
    def from_samples(
        cls,
        log_prob_fn: Callable[[mx.array, mx.array], mx.array],
        theta: mx.array,
        samples: mx.array,
        eps: float = 1e-4,
    ) -> "FisherMatrix":
        """
        Empirical Fisher from samples:

            F_ij = (1/N) sum_n  score_i(x_n; theta) * score_j(x_n; theta)

        Parameters
        ----------
        log_prob_fn : callable (theta, x) -> scalar
            Log-probability of a single sample x given parameters theta.
        theta : mx.array, shape (n_params,)
        samples : mx.array, shape (n_samples, ...)
        eps : float
            Step for numerical gradient of log_prob w.r.t. theta.

        Returns
        -------
        FisherMatrix
        """
        theta = mx.array(theta, dtype=mx.float32)
        n_params = theta.shape[0]
        n_samples = samples.shape[0]

        # Compute score vectors for every sample
        all_scores = []
        for n in range(n_samples):
            x_n = samples[n]
            score_n = mx.zeros((n_params,), dtype=mx.float32)
            scores_i = []
            for i in range(n_params):
                e_i = mx.zeros_like(theta)
                e_i[i] = eps
                lp_plus = log_prob_fn(theta + e_i, x_n)
                lp_minus = log_prob_fn(theta - e_i, x_n)
                scores_i.append((lp_plus - lp_minus) / (2.0 * eps))
            all_scores.append(mx.stack(scores_i))

        # S: (n_samples, n_params)
        S = mx.stack(all_scores, axis=0)
        F = (S.T @ S) / n_samples
        mx.eval(F)
        return cls(F)

    # ------------------------------------------------------------------
    # Algebra
    # ------------------------------------------------------------------

    def __add__(self, other: "FisherMatrix") -> "FisherMatrix":
        """Add two Fisher matrices (e.g. combining independent experiments)."""
        return FisherMatrix(self._matrix + other._matrix)

    def __repr__(self) -> str:
        return f"FisherMatrix(n_params={self.n_params}, cond={self.condition_number():.2e})"


# ======================================================================
# Specialized: CMB C_l Fisher matrix
# ======================================================================


def fisher_matrix_cl(
    cl_fn: Callable[[mx.array], mx.array],
    theta: mx.array,
    f_sky: float = 1.0,
    noise_cl: Optional[mx.array] = None,
    l_min: int = 2,
    l_max: Optional[int] = None,
    eps: float = 1e-4,
) -> FisherMatrix:
    """
    Fisher matrix for CMB power spectrum C_l.

    For a Gaussian, isotropic sky the Fisher matrix is:

        F_ij = sum_{l=l_min}^{l_max} (2l+1) f_sky / (2 C_l_tot^2)
               * (d C_l / d theta_i) * (d C_l / d theta_j)

    where C_l_tot = C_l(theta) + N_l.

    Parameters
    ----------
    cl_fn : callable  theta -> mx.array (n_ell,)
        Returns C_l starting at l=0. The first l_min entries are ignored.
    theta : mx.array, shape (n_params,)
        Fiducial cosmological parameters.
    f_sky : float
        Sky fraction (0 < f_sky <= 1).
    noise_cl : mx.array or None, shape (n_ell,)
        Noise power spectrum N_l.  If None, assume zero noise.
    l_min : int
        Minimum multipole (default 2).
    l_max : int or None
        Maximum multipole. If None, use full length of cl_fn output.
    eps : float
        Step size for numerical derivatives.

    Returns
    -------
    FisherMatrix
    """
    theta = mx.array(theta, dtype=mx.float32)
    n_params = theta.shape[0]

    # Fiducial C_l
    cl_fid = cl_fn(theta)
    n_ell = cl_fid.shape[0]
    if l_max is None:
        l_max = n_ell - 1

    # Noise
    if noise_cl is None:
        noise_cl = mx.zeros((n_ell,), dtype=mx.float32)
    else:
        noise_cl = mx.array(noise_cl, dtype=mx.float32)

    # Total C_l
    cl_tot = cl_fid + noise_cl

    # Multipole weights: (2l+1) * f_sky / (2 * C_l_tot^2)
    ell = mx.arange(n_ell, dtype=mx.float32)
    weights = (2.0 * ell + 1.0) * f_sky / (2.0 * cl_tot * cl_tot)

    # Zero out l < l_min and l > l_max
    mask = mx.zeros((n_ell,), dtype=mx.float32)
    indices = mx.arange(n_ell)
    mask = mx.where((indices >= l_min) & (indices <= l_max), 1.0, 0.0)
    weights = weights * mask

    # Derivatives d C_l / d theta_i via central differences
    dcl = []
    for i in range(n_params):
        e_i = mx.zeros_like(theta)
        e_i[i] = eps
        cl_plus = cl_fn(theta + e_i)
        cl_minus = cl_fn(theta - e_i)
        dcl.append((cl_plus - cl_minus) / (2.0 * eps))

    # Assemble Fisher matrix
    F = mx.zeros((n_params, n_params), dtype=mx.float32)
    for i in range(n_params):
        for j in range(i, n_params):
            F_ij = mx.sum(weights * dcl[i] * dcl[j])
            F[i, j] = F_ij
            F[j, i] = F_ij

    mx.eval(F)
    return FisherMatrix(F)


# ======================================================================
# KL divergence (classical)
# ======================================================================


def kl_divergence(p: mx.array, q: mx.array, eps: float = 1e-30) -> mx.array:
    """
    Kullback-Leibler divergence D_KL(p || q) for discrete distributions.

        D_KL = sum_i p_i * ln(p_i / q_i)

    Parameters
    ----------
    p, q : mx.array, shape (n,)
        Probability distributions (must sum to 1, all entries >= 0).
    eps : float
        Small constant to avoid log(0).

    Returns
    -------
    mx.array (scalar)
    """
    p = mx.clip(p, eps, None)
    q = mx.clip(q, eps, None)
    return mx.sum(p * (mx.log(p) - mx.log(q)))


# ======================================================================
# Quantum relative entropy
# ======================================================================


def quantum_relative_entropy(rho: mx.array, sigma: mx.array, eps: float = 1e-12) -> mx.array:
    """
    Quantum relative entropy (Umegaki):

        D(rho || sigma) = Tr[ rho (ln rho - ln sigma) ]

    This is the Sigma formula.

    Parameters
    ----------
    rho : mx.array, shape (d, d)
        Density matrix (Hermitian, positive semi-definite, Tr=1).
    sigma : mx.array, shape (d, d)
        Density matrix.
    eps : float
        Regularisation for eigenvalues to avoid log(0).

    Returns
    -------
    mx.array (scalar)
        The quantum relative entropy D(rho || sigma).
    """
    d = rho.shape[0]

    # Diagonalise rho
    eigvals_rho, U_rho = mx.linalg.eigh(rho, stream=mx.cpu)
    eigvals_rho = mx.clip(eigvals_rho, eps, None)
    log_rho = U_rho @ mx.diag(mx.log(eigvals_rho)) @ U_rho.T

    # Diagonalise sigma
    eigvals_sigma, U_sigma = mx.linalg.eigh(sigma, stream=mx.cpu)
    eigvals_sigma = mx.clip(eigvals_sigma, eps, None)
    log_sigma = U_sigma @ mx.diag(mx.log(eigvals_sigma)) @ U_sigma.T

    # D(rho || sigma) = Tr[ rho (log_rho - log_sigma) ]
    diff = log_rho - log_sigma
    result = _trace_product(rho, diff)
    mx.eval(result)
    return result


def _trace_product(A: mx.array, B: mx.array) -> mx.array:
    """Tr(A @ B) computed efficiently as sum of element-wise product."""
    return mx.sum(A * B.T)
