"""
Benchmark: MLX (GPU) vs NumPy (CPU) Fisher matrix computation.

Compares wall-clock time for:
  1. Gaussian Fisher matrix (model Jacobian method)
  2. Matrix eigendecomposition
  3. KL divergence
  4. Quantum relative entropy

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

from __future__ import annotations

import time
from dataclasses import dataclass

import mlx.core as mx
import numpy as np


@dataclass
class BenchmarkResult:
    name: str
    n_params: int
    n_data: int
    mlx_time_ms: float
    numpy_time_ms: float

    @property
    def speedup(self) -> float:
        if self.mlx_time_ms <= 0:
            return float("inf")
        return self.numpy_time_ms / self.mlx_time_ms


def _warmup_mlx():
    """Warm up the MLX GPU by running a small computation."""
    a = mx.ones((64, 64))
    b = a @ a
    mx.eval(b)


# ======================================================================
# Benchmark 1: Gaussian Fisher matrix
# ======================================================================


def benchmark_gaussian_fisher(
    n_params: int = 6, n_data: int = 2500, n_repeats: int = 5
) -> BenchmarkResult:
    """
    Benchmark Gaussian Fisher matrix: F_ij = J^T diag(1/sigma^2) J
    where J is (n_data, n_params) Jacobian.
    """
    # --- MLX ---
    _warmup_mlx()
    J_mx = mx.random.normal((n_data, n_params))
    sigma_mx = mx.abs(mx.random.normal((n_data,))) + 0.1
    mx.eval(J_mx, sigma_mx)

    # Warm-up run
    inv_var = 1.0 / (sigma_mx * sigma_mx)
    J_w = J_mx * inv_var[:, None]
    F = J_w.T @ J_mx
    mx.eval(F)

    t0 = time.perf_counter()
    for _ in range(n_repeats):
        inv_var = 1.0 / (sigma_mx * sigma_mx)
        J_w = J_mx * inv_var[:, None]
        F = J_w.T @ J_mx
        mx.eval(F)
    mlx_time = (time.perf_counter() - t0) / n_repeats

    # --- NumPy ---
    J_np = np.array(J_mx)
    sigma_np = np.array(sigma_mx)

    t0 = time.perf_counter()
    for _ in range(n_repeats):
        inv_var_np = 1.0 / (sigma_np ** 2)
        J_w_np = J_np * inv_var_np[:, None]
        F_np = J_w_np.T @ J_np
    numpy_time = (time.perf_counter() - t0) / n_repeats

    return BenchmarkResult(
        name="Gaussian Fisher (J^T W J)",
        n_params=n_params,
        n_data=n_data,
        mlx_time_ms=mlx_time * 1000,
        numpy_time_ms=numpy_time * 1000,
    )


# ======================================================================
# Benchmark 2: Eigendecomposition
# ======================================================================


def benchmark_eigendecomposition(
    n: int = 512, n_repeats: int = 5
) -> BenchmarkResult:
    """Benchmark symmetric eigendecomposition of an n x n matrix."""
    _warmup_mlx()

    # Create a symmetric positive-definite matrix
    A_mx = mx.random.normal((n, n))
    A_mx = A_mx @ A_mx.T + mx.eye(n)
    mx.eval(A_mx)

    # Warm-up
    eigvals, eigvecs = mx.linalg.eigh(A_mx, stream=mx.cpu)
    mx.eval(eigvals, eigvecs)

    t0 = time.perf_counter()
    for _ in range(n_repeats):
        eigvals, eigvecs = mx.linalg.eigh(A_mx, stream=mx.cpu)
        mx.eval(eigvals, eigvecs)
    mlx_time = (time.perf_counter() - t0) / n_repeats

    # NumPy
    A_np = np.array(A_mx)
    t0 = time.perf_counter()
    for _ in range(n_repeats):
        eigvals_np, eigvecs_np = np.linalg.eigh(A_np)
    numpy_time = (time.perf_counter() - t0) / n_repeats

    return BenchmarkResult(
        name="Eigendecomposition (eigh)",
        n_params=n,
        n_data=n,
        mlx_time_ms=mlx_time * 1000,
        numpy_time_ms=numpy_time * 1000,
    )


# ======================================================================
# Benchmark 3: KL divergence
# ======================================================================


def benchmark_kl_divergence(
    n: int = 100_000, n_repeats: int = 10
) -> BenchmarkResult:
    """Benchmark discrete KL divergence for distributions of size n."""
    _warmup_mlx()

    # Random distributions
    p_mx = mx.abs(mx.random.normal((n,))) + 1e-6
    p_mx = p_mx / mx.sum(p_mx)
    q_mx = mx.abs(mx.random.normal((n,))) + 1e-6
    q_mx = q_mx / mx.sum(q_mx)
    mx.eval(p_mx, q_mx)

    # Warm-up
    kl = mx.sum(p_mx * (mx.log(p_mx) - mx.log(q_mx)))
    mx.eval(kl)

    t0 = time.perf_counter()
    for _ in range(n_repeats):
        kl = mx.sum(p_mx * (mx.log(p_mx) - mx.log(q_mx)))
        mx.eval(kl)
    mlx_time = (time.perf_counter() - t0) / n_repeats

    # NumPy
    p_np = np.array(p_mx)
    q_np = np.array(q_mx)
    t0 = time.perf_counter()
    for _ in range(n_repeats):
        kl_np = np.sum(p_np * (np.log(p_np) - np.log(q_np)))
    numpy_time = (time.perf_counter() - t0) / n_repeats

    return BenchmarkResult(
        name="KL divergence",
        n_params=1,
        n_data=n,
        mlx_time_ms=mlx_time * 1000,
        numpy_time_ms=numpy_time * 1000,
    )


# ======================================================================
# Benchmark 4: Quantum relative entropy
# ======================================================================


def benchmark_quantum_relative_entropy(
    d: int = 256, n_repeats: int = 5
) -> BenchmarkResult:
    """
    Benchmark quantum relative entropy D(rho || sigma) for d x d density matrices.
    Involves two eigendecompositions + matrix multiplies + trace.
    """
    _warmup_mlx()
    eps = 1e-12

    # Random density matrices
    def _random_density_mx(dim):
        A = mx.random.normal((dim, dim))
        rho = A @ A.T
        rho = rho / mx.trace(rho)
        mx.eval(rho)
        return rho

    def _random_density_np(dim):
        A = np.random.randn(dim, dim).astype(np.float32)
        rho = A @ A.T
        rho = rho / np.trace(rho)
        return rho

    rho_mx = _random_density_mx(d)
    sigma_mx = _random_density_mx(d)

    # MLX warm-up
    eigr, Ur = mx.linalg.eigh(rho_mx, stream=mx.cpu)
    mx.eval(eigr, Ur)

    t0 = time.perf_counter()
    for _ in range(n_repeats):
        eigr, Ur = mx.linalg.eigh(rho_mx, stream=mx.cpu)
        eigr = mx.clip(eigr, eps, None)
        log_rho = Ur @ mx.diag(mx.log(eigr)) @ Ur.T

        eigs, Us = mx.linalg.eigh(sigma_mx, stream=mx.cpu)
        eigs = mx.clip(eigs, eps, None)
        log_sigma = Us @ mx.diag(mx.log(eigs)) @ Us.T

        D = mx.sum(rho_mx * (log_rho - log_sigma).T)
        mx.eval(D)
    mlx_time = (time.perf_counter() - t0) / n_repeats

    # NumPy
    rho_np = np.array(rho_mx)
    sigma_np = np.array(sigma_mx)

    t0 = time.perf_counter()
    for _ in range(n_repeats):
        eigr_np, Ur_np = np.linalg.eigh(rho_np)
        eigr_np = np.clip(eigr_np, eps, None)
        log_rho_np = Ur_np @ np.diag(np.log(eigr_np)) @ Ur_np.T

        eigs_np, Us_np = np.linalg.eigh(sigma_np)
        eigs_np = np.clip(eigs_np, eps, None)
        log_sigma_np = Us_np @ np.diag(np.log(eigs_np)) @ Us_np.T

        D_np = np.sum(rho_np * (log_rho_np - log_sigma_np).T)
    numpy_time = (time.perf_counter() - t0) / n_repeats

    return BenchmarkResult(
        name="Quantum relative entropy",
        n_params=d,
        n_data=d,
        mlx_time_ms=mlx_time * 1000,
        numpy_time_ms=numpy_time * 1000,
    )


# ======================================================================
# Benchmark 5: Large matrix multiply (Fisher-like)
# ======================================================================


def benchmark_matmul(
    m: int = 2048, n: int = 64, n_repeats: int = 5
) -> BenchmarkResult:
    """Benchmark J^T @ J where J is (m, n) -- typical Fisher assembly."""
    _warmup_mlx()

    J_mx = mx.random.normal((m, n))
    mx.eval(J_mx)

    # Warm-up
    F = J_mx.T @ J_mx
    mx.eval(F)

    t0 = time.perf_counter()
    for _ in range(n_repeats):
        F = J_mx.T @ J_mx
        mx.eval(F)
    mlx_time = (time.perf_counter() - t0) / n_repeats

    J_np = np.array(J_mx)
    t0 = time.perf_counter()
    for _ in range(n_repeats):
        F_np = J_np.T @ J_np
    numpy_time = (time.perf_counter() - t0) / n_repeats

    return BenchmarkResult(
        name=f"Matrix multiply J^T J ({m}x{n})",
        n_params=n,
        n_data=m,
        mlx_time_ms=mlx_time * 1000,
        numpy_time_ms=numpy_time * 1000,
    )


# ======================================================================
# Run all benchmarks
# ======================================================================


def run_all(verbose: bool = True) -> list[BenchmarkResult]:
    """Run all benchmarks and return results."""
    benchmarks = [
        ("Gaussian Fisher (6 params, 2500 data)", lambda: benchmark_gaussian_fisher(6, 2500)),
        ("Gaussian Fisher (10 params, 10000 data)", lambda: benchmark_gaussian_fisher(10, 10000)),
        ("Eigendecomposition 256x256", lambda: benchmark_eigendecomposition(256)),
        ("Eigendecomposition 512x512", lambda: benchmark_eigendecomposition(512)),
        ("KL divergence (100k bins)", lambda: benchmark_kl_divergence(100_000)),
        ("KL divergence (1M bins)", lambda: benchmark_kl_divergence(1_000_000)),
        ("Quantum relative entropy (128x128)", lambda: benchmark_quantum_relative_entropy(128)),
        ("Quantum relative entropy (256x256)", lambda: benchmark_quantum_relative_entropy(256)),
        ("Matrix multiply (2048x64)", lambda: benchmark_matmul(2048, 64)),
        ("Matrix multiply (8192x128)", lambda: benchmark_matmul(8192, 128)),
    ]

    results = []
    for name, fn in benchmarks:
        if verbose:
            print(f"Running: {name} ... ", end="", flush=True)
        result = fn()
        results.append(result)
        if verbose:
            print(f"MLX {result.mlx_time_ms:.3f} ms  |  NumPy {result.numpy_time_ms:.3f} ms  |  speedup {result.speedup:.2f}x")

    return results


def format_markdown(results: list[BenchmarkResult]) -> str:
    """Format benchmark results as a Markdown table."""
    lines = [
        "| Benchmark | Dimensions | MLX (ms) | NumPy (ms) | Speedup |",
        "|-----------|-----------|----------|------------|---------|",
    ]
    for r in results:
        dims = f"{r.n_data} x {r.n_params}" if r.n_data != r.n_params else f"{r.n_params} x {r.n_params}"
        lines.append(
            f"| {r.name} | {dims} | {r.mlx_time_ms:.3f} | {r.numpy_time_ms:.3f} | {r.speedup:.2f}x |"
        )
    return "\n".join(lines)


if __name__ == "__main__":
    print("=" * 70)
    print("mlx-fisher benchmark: MLX (GPU) vs NumPy (CPU)")
    print("=" * 70)
    print()
    results = run_all(verbose=True)
    print()
    print("Markdown table:")
    print()
    print(format_markdown(results))
