"""Quadratic/Gaussian energy model for BABAPPA."""

from __future__ import annotations

from dataclasses import asdict, dataclass

import numpy as np


@dataclass
class GaussianEnergyModel:
    mean: list[float]
    precision: list[list[float]]
    ridge: float
    n_training_sites: int
    score_matching_objective: float
    convergence_flag: str
    min_eigenvalue: float
    max_eigenvalue: float

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


@dataclass
class FittedGaussian:
    model: GaussianEnergyModel
    mean_vec: np.ndarray
    precision_mat: np.ndarray


def _score_matching_objective(samples: np.ndarray, precision: np.ndarray) -> float:
    # Hyvarinen score objective for Gaussian (up to constants):
    # J(K) = 0.5 * E[||K(x-mu)||^2] - tr(K)
    centered = samples - samples.mean(axis=0, keepdims=True)
    quad = np.einsum("ni,ij,nj->n", centered, precision @ precision, centered)
    return float(0.5 * np.mean(quad) - np.trace(precision))


def fit_gaussian_energy(samples: np.ndarray, ridge: float = 1e-4) -> FittedGaussian:
    if samples.ndim != 2:
        raise ValueError("Samples must be a 2D array [n_sites, dim]")
    if samples.shape[0] < 3:
        raise ValueError("At least 3 training sites are required")

    mu = samples.mean(axis=0)
    centered = samples - mu
    cov = centered.T @ centered / max(1, samples.shape[0] - 1)
    cov = cov + ridge * np.eye(cov.shape[0])

    eigvals = np.linalg.eigvalsh(cov)
    precision = np.linalg.pinv(cov, hermitian=True)

    objective = _score_matching_objective(samples, precision)
    model = GaussianEnergyModel(
        mean=mu.tolist(),
        precision=precision.tolist(),
        ridge=float(ridge),
        n_training_sites=int(samples.shape[0]),
        score_matching_objective=float(objective),
        convergence_flag="closed_form",
        min_eigenvalue=float(np.min(eigvals)),
        max_eigenvalue=float(np.max(eigvals)),
    )
    return FittedGaussian(model=model, mean_vec=mu, precision_mat=precision)


def site_mahalanobis_scores(site_vectors: np.ndarray, fitted: FittedGaussian) -> np.ndarray:
    centered = site_vectors - fitted.mean_vec
    return np.einsum("ni,ij,nj->n", centered, fitted.precision_mat, centered)


def gene_dispersion_statistic(site_vectors: np.ndarray, fitted: FittedGaussian) -> tuple[float, np.ndarray]:
    per_site = site_mahalanobis_scores(site_vectors, fitted=fitted)
    if per_site.size == 0:
        raise ValueError("At least one site is required to compute gene dispersion")
    mean_energy = float(np.mean(per_site))
    dispersion = float(np.mean((per_site - mean_energy) ** 2))
    return dispersion, per_site


def fitted_from_dict(payload: dict[str, object]) -> FittedGaussian:
    model = GaussianEnergyModel(**payload)
    mean_vec = np.array(model.mean, dtype=float)
    precision = np.array(model.precision, dtype=float)
    return FittedGaussian(model=model, mean_vec=mean_vec, precision_mat=precision)
