#pragma once

#include "spectra_reml/types.hpp"

#include <limits>
#include <string>
#include <vector>

namespace spectra::reml {

// Fits y = X beta + g + e in the GRM eigenspace, where
// Var(y) = sigma_g^2 diag(lambda) + sigma_e^2 I.
//
// The parameters are the signed standard deviations (sigma_e, sigma_g), as in
// thesis Equations 4.24-4.25.  The covariance depends on their squares, so the
// optimization is unconstrained.  It uses AI-REML and a strong-Wolfe line
// search; no EM update is performed and P is never materialized.
[[nodiscard]] RemlResult fit_ai_reml_spectral(
    const std::vector<double>& y_star, const ColMajorMatrix& x_star,
    const std::vector<double>& eigenvalues,
    const RemlOptions& options = {});

// Exposed for finite-difference and independent-oracle tests.  The ordering is
// theta=(sigma_e,sigma_g); ai is [ee,eg;eg,gg] in column-major order.
struct RemlEvaluation {
    bool valid = false;
    double log_likelihood = -std::numeric_limits<double>::infinity();
    double gradient_e = std::numeric_limits<double>::quiet_NaN();
    double gradient_g = std::numeric_limits<double>::quiet_NaN();
    // Scores with respect to the variance parameters v_e=sigma_e^2 and
    // v_g=sigma_g^2.  These remain informative when sigma_g=0, unlike the
    // signed-standard-deviation gradient gradient_g=2*sigma_g*score_v_g.
    double variance_score_e = std::numeric_limits<double>::quiet_NaN();
    double variance_score_g = std::numeric_limits<double>::quiet_NaN();
    double ai_ee = std::numeric_limits<double>::quiet_NaN();
    double ai_eg = std::numeric_limits<double>::quiet_NaN();
    double ai_gg = std::numeric_limits<double>::quiet_NaN();
    std::vector<double> beta;
    std::vector<double> beta_covariance;
    std::string error;
};

[[nodiscard]] RemlEvaluation evaluate_reml_spectral(
    const std::vector<double>& y_star, const ColMajorMatrix& x_star,
    const std::vector<double>& eigenvalues, double sigma_e, double sigma_g,
    bool compute_ai = true, bool compute_beta_covariance = false,
    double covariance_floor_relative = 1e-12);

}  // namespace spectra::reml