#include "spectra_reml/reml.hpp"
#include "spectra_reml/fixed_effects.hpp"
#include "spectra_reml/distributions.hpp"

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

void require(bool condition, const std::string& message) {
    if (!condition) {
        throw std::runtime_error(message);
    }
}

void require_near(double actual, double expected, double relative_tolerance,
                  const std::string& message) {
    const double scale = std::max({1.0, std::abs(actual), std::abs(expected)});
    if (!std::isfinite(actual) || !std::isfinite(expected) ||
        std::abs(actual - expected) > relative_tolerance * scale) {
        throw std::runtime_error(message + ": actual=" +
                                 std::to_string(actual) +
                                 ", expected=" + std::to_string(expected));
    }
}

struct Fixture {
    std::vector<double> y;
    spectra::reml::ColMajorMatrix x;
    std::vector<double> lambda;

    Fixture() : x(18, 3) {
        lambda = {0.12, 0.18, 0.25, 0.31, 0.39, 0.48, 0.58, 0.69, 0.81,
                  0.95, 1.10, 1.28, 1.47, 1.69, 1.94, 2.22, 2.55, 2.91};
        y.resize(lambda.size());
        for (std::size_t i = 0; i < lambda.size(); ++i) {
            const double z = (static_cast<double>(i) - 8.5) / 5.0;
            x(i, 0) = 1.0;
            x(i, 1) = z;
            x(i, 2) = (i % 3 == 0 ? -0.8 : (i % 3 == 1 ? 0.1 : 0.7));
            const double deterministic_noise =
                std::sin(1.7 * static_cast<double>(i) + 0.3) *
                    std::sqrt(0.35 + 0.55 * lambda[i]) +
                0.22 * std::cos(0.37 * static_cast<double>(i));
            y[i] = 0.4 - 0.25 * x(i, 1) + 0.15 * x(i, 2) +
                   deterministic_noise;
        }
    }
};

void test_score_matches_finite_difference() {
    Fixture fixture;
    constexpr double sigma_e = 0.73;
    constexpr double sigma_g = 0.46;
    const auto evaluated = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g, true, true);
    require(evaluated.valid, "reference REML evaluation failed: " +
                                 evaluated.error);
    require(evaluated.beta.size() == fixture.x.cols(),
            "beta has wrong length");
    require(evaluated.beta_covariance.size() ==
                fixture.x.cols() * fixture.x.cols(),
            "beta covariance has wrong size");

    const double epsilon_e = 2e-6 * std::max(1.0, std::abs(sigma_e));
    const double epsilon_g = 2e-6 * std::max(1.0, std::abs(sigma_g));
    const auto e_plus = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, sigma_e + epsilon_e, sigma_g,
        false, false);
    const auto e_minus = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, sigma_e - epsilon_e, sigma_g,
        false, false);
    const auto g_plus = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g + epsilon_g,
        false, false);
    const auto g_minus = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g - epsilon_g,
        false, false);
    require(e_plus.valid && e_minus.valid && g_plus.valid && g_minus.valid,
            "finite-difference REML evaluation failed");
    const double finite_difference_e =
        (e_plus.log_likelihood - e_minus.log_likelihood) / (2.0 * epsilon_e);
    const double finite_difference_g =
        (g_plus.log_likelihood - g_minus.log_likelihood) / (2.0 * epsilon_g);
    require_near(evaluated.gradient_e, finite_difference_e, 2e-6,
                 "sigma_e score disagrees with finite difference");
    require_near(evaluated.gradient_g, finite_difference_g, 2e-6,
                 "sigma_g score disagrees with finite difference");
    require(evaluated.ai_ee > 0.0 && evaluated.ai_gg > 0.0,
            "AI diagonal must be positive");
    require(evaluated.ai_ee * evaluated.ai_gg -
                    evaluated.ai_eg * evaluated.ai_eg >=
                -1e-10,
            "AI matrix must be positive semidefinite");
}

void test_signed_parameterization_and_row_permutation() {
    Fixture fixture;
    const auto positive = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, 0.61, 0.52, true, false);
    const auto negative = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, -0.61, -0.52, true, false);
    require(positive.valid && negative.valid, "signed evaluations failed");
    require_near(positive.log_likelihood, negative.log_likelihood, 1e-13,
                 "likelihood must depend on squared standard deviations");
    require_near(positive.gradient_e, -negative.gradient_e, 1e-12,
                 "sigma_e gradient must change sign");
    require_near(positive.gradient_g, -negative.gradient_g, 1e-12,
                 "sigma_g gradient must change sign");

    std::vector<double> permuted_y(fixture.y.rbegin(), fixture.y.rend());
    std::vector<double> permuted_lambda(fixture.lambda.rbegin(),
                                        fixture.lambda.rend());
    spectra::reml::ColMajorMatrix permuted_x(fixture.x.rows(), fixture.x.cols());
    for (std::size_t i = 0; i < fixture.x.rows(); ++i) {
        for (std::size_t j = 0; j < fixture.x.cols(); ++j) {
            permuted_x(i, j) = fixture.x(fixture.x.rows() - 1 - i, j);
        }
    }
    const auto permuted = spectra::reml::evaluate_reml_spectral(
        permuted_y, permuted_x, permuted_lambda, 0.61, 0.52, true, false);
    require(permuted.valid, "permuted evaluation failed");
    require_near(positive.log_likelihood, permuted.log_likelihood, 2e-13,
                 "orthogonal-axis permutation changed REML likelihood");
    require_near(positive.gradient_e, permuted.gradient_e, 2e-12,
                 "orthogonal-axis permutation changed sigma_e score");
    require_near(positive.gradient_g, permuted.gradient_g, 2e-12,
                 "orthogonal-axis permutation changed sigma_g score");
}

void test_ai_reml_fit_improves_likelihood() {
    Fixture fixture;
    spectra::reml::RemlOptions options;
    options.initial_sigma_e = 0.9;
    options.initial_sigma_g = 0.9;
    options.max_iterations = 150;
    options.gradient_absolute_tolerance = 1e-8;
    options.gradient_relative_tolerance = 1e-9;
    const auto initial = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, options.initial_sigma_e,
        options.initial_sigma_g, true, false);
    require(initial.valid, "initial fit evaluation failed");
    const auto fitted = spectra::reml::fit_ai_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, options);
    require(fitted.has_estimates(),
            "AI-REML returned no estimates: " + fitted.error);
    require(fitted.status == spectra::reml::FitStatus::converged,
            std::string("AI-REML did not converge: ") +
                spectra::reml::to_string(fitted.status) + " " + fitted.error);
    require(fitted.log_likelihood >= initial.log_likelihood - 1e-10,
            "AI-REML decreased the likelihood");
    require(fitted.gradient_inf_norm < 2e-5,
            "AI-REML final score is too large");
    require(fitted.beta.size() == fixture.x.cols(),
            "fitted beta has wrong length");
    require(fitted.beta_covariance_packed_lower.size() ==
                fixture.x.cols() * (fixture.x.cols() + 1) / 2,
            "packed beta covariance has wrong length");
    require_near(fitted.sigma_e2, fitted.sigma_e * fitted.sigma_e, 1e-14,
                 "reported sigma_e2 is inconsistent");
    require_near(fitted.sigma_g2, fitted.sigma_g * fitted.sigma_g, 1e-14,
                 "reported sigma_g2 is inconsistent");

    spectra::reml::FixedEffectHypothesis joint;
    joint.contrast = spectra::reml::ColMajorMatrix(2, fixture.x.cols());
    joint.contrast(0, 1) = 1.0;
    joint.contrast(1, 2) = 1.0;
    const auto satterthwaite = spectra::reml::infer_fixed_effects_spectral(
        fixture.y, fixture.x, fixture.lambda, fitted,
        spectra::reml::FixedEffectTestMethod::satterthwaite, {joint});
    require(satterthwaite.status ==
                spectra::reml::FixedEffectInferenceStatus::ok,
            "interior Satterthwaite inference failed: " + satterthwaite.error);
    require(satterthwaite.coefficient_tests.size() == fixture.x.cols(),
            "interior Satterthwaite coefficient tests are missing");
    require(satterthwaite.hypothesis_tests.size() == 1 &&
                satterthwaite.hypothesis_tests.front().numerator_df == 2,
            "interior Satterthwaite joint test is missing");
    require(satterthwaite.hypothesis_tests.front().valid &&
                satterthwaite.hypothesis_tests.front().p_value >= 0.0 &&
                satterthwaite.hypothesis_tests.front().p_value <= 1.0,
            "interior Satterthwaite joint p-value is invalid");

    const auto kr = spectra::reml::infer_fixed_effects_spectral(
        fixture.y, fixture.x, fixture.lambda, fitted,
        spectra::reml::FixedEffectTestMethod::kenward_roger, {joint});
    require(kr.status == spectra::reml::FixedEffectInferenceStatus::ok,
            "interior KR inference failed: " + kr.error);
    require(kr.hypothesis_tests.size() == 1 &&
                kr.hypothesis_tests.front().valid &&
                kr.hypothesis_tests.front().numerator_df == 2,
            "interior KR joint test is invalid");

    auto line_search_flagged = fitted;
    line_search_flagged.status =
        spectra::reml::FitStatus::line_search_failed;
    line_search_flagged.error = "synthetic retained line-search iterate";
    const auto flagged_satterthwaite =
        spectra::reml::infer_fixed_effects_spectral(
            fixture.y, fixture.x, fixture.lambda, line_search_flagged,
            spectra::reml::FixedEffectTestMethod::satterthwaite, {joint});
    require(flagged_satterthwaite.status ==
                spectra::reml::FixedEffectInferenceStatus::line_search_conditional,
            "line-search Satterthwaite inference was not retained and flagged");
    require(flagged_satterthwaite.hypothesis_tests.size() == 1 &&
                flagged_satterthwaite.hypothesis_tests.front().valid,
            "line-search Satterthwaite hypothesis test is invalid");
    const auto flagged_kr = spectra::reml::infer_fixed_effects_spectral(
        fixture.y, fixture.x, fixture.lambda, line_search_flagged,
        spectra::reml::FixedEffectTestMethod::kenward_roger, {joint});
    require(flagged_kr.status ==
                spectra::reml::FixedEffectInferenceStatus::line_search_conditional,
            "line-search KR inference was not retained and flagged");
    require(flagged_kr.hypothesis_tests.size() == 1 &&
                flagged_kr.hypothesis_tests.front().valid,
            "line-search KR hypothesis test is invalid");
}

void test_probability_distributions_against_r() {
    require_near(spectra::reml::student_t_two_sided_p(2.1, 7.3),
                 0.072246713424853351, 2e-14,
                 "Student-t tail probability disagrees with R");
    require_near(spectra::reml::f_upper_tail(3.7, 2.5, 11.2),
                 0.051114185450108554, 2e-14,
                 "F tail probability disagrees with R");
    require_near(spectra::reml::student_t_two_sided_p(12.0, 3.5),
                 0.0005745341928561279, 2e-14,
                 "extreme Student-t tail probability disagrees with R");
    require_near(spectra::reml::f_upper_tail(100.0, 1.0, 8.25),
                 6.803231179451616e-06, 2e-14,
                 "extreme F tail probability disagrees with R");
}

void test_fixed_effect_inference_against_dense_reference() {
    Fixture fixture;
    constexpr double sigma_e = 0.73;
    constexpr double sigma_g = 0.46;
    const auto evaluated = spectra::reml::evaluate_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g, false, true);
    require(evaluated.valid, "fixed-effect reference evaluation failed");
    spectra::reml::RemlResult fit;
    fit.status = spectra::reml::FitStatus::converged;
    fit.sigma_e = sigma_e;
    fit.sigma_g = sigma_g;
    fit.sigma_e2 = sigma_e * sigma_e;
    fit.sigma_g2 = sigma_g * sigma_g;
    fit.beta = evaluated.beta;

    spectra::reml::FixedEffectHypothesis joint;
    joint.contrast = spectra::reml::ColMajorMatrix(2, fixture.x.cols());
    joint.contrast(0, 1) = 1.0;
    joint.contrast(1, 2) = 1.0;
    const auto satt = spectra::reml::infer_fixed_effects_spectral(
        fixture.y, fixture.x, fixture.lambda, fit,
        spectra::reml::FixedEffectTestMethod::satterthwaite, {joint});
    require(satt.status == spectra::reml::FixedEffectInferenceStatus::ok,
            "dense-reference Satterthwaite inference failed: " + satt.error);
    const std::vector<double> expected_satt_df = {
        11.271857730561006, 11.012661088553992, 8.7818755201562855};
    const std::vector<double> expected_satt_t = {
        1.9994269321713118, -1.6229945279594145, 0.26284041305853334};
    for (std::size_t i = 0; i < fixture.x.cols(); ++i) {
        require_near(satt.coefficient_tests[i].denominator_df,
                     expected_satt_df[i], 3e-11,
                     "Satterthwaite df disagrees with dense reference");
        require_near(satt.coefficient_tests[i].statistic,
                     expected_satt_t[i], 3e-12,
                     "Satterthwaite t disagrees with dense reference");
    }
    require_near(satt.hypothesis_tests[0].statistic,
                 1.3172663846541173, 3e-12,
                 "Satterthwaite joint F disagrees with dense reference");
    require_near(satt.hypothesis_tests[0].denominator_df,
                 9.7667759456518599, 3e-11,
                 "Satterthwaite joint df disagrees with dense reference");

    const auto kr = spectra::reml::infer_fixed_effects_spectral(
        fixture.y, fixture.x, fixture.lambda, fit,
        spectra::reml::FixedEffectTestMethod::kenward_roger, {joint});
    require(kr.status == spectra::reml::FixedEffectInferenceStatus::ok,
            "dense-reference KR inference failed: " + kr.error);
    const std::vector<double> expected_kr_df = {
        14.488430939129156, 14.905537974410061, 14.020384661469432};
    const std::vector<double> expected_kr_f = {
        3.9170644186264694, 2.1360773364244796, 0.0544837858871276};
    for (std::size_t i = 0; i < fixture.x.cols(); ++i) {
        require_near(kr.coefficient_tests[i].denominator_df,
                     expected_kr_df[i], 3e-11,
                     "KR df disagrees with dense reference");
        require_near(kr.coefficient_tests[i].statistic,
                     expected_kr_f[i], 3e-11,
                     "KR F disagrees with dense reference");
    }
    require_near(kr.hypothesis_tests[0].statistic,
                 1.0684514230854425, 3e-11,
                 "KR joint F disagrees with dense reference");
    require_near(kr.hypothesis_tests[0].denominator_df,
                 14.615505880683108, 3e-11,
                 "KR joint df disagrees with dense reference");
}

void test_internal_phenotype_scaling_restores_original_units() {
    Fixture fixture;
    constexpr double scale = 1.0e6;
    std::vector<double> scaled_y = fixture.y;
    for (double& value : scaled_y) {
        value *= scale;
    }
    spectra::reml::RemlOptions options;
    options.max_iterations = 150;
    const auto original = spectra::reml::fit_ai_reml_spectral(
        fixture.y, fixture.x, fixture.lambda, options);
    const auto scaled = spectra::reml::fit_ai_reml_spectral(
        scaled_y, fixture.x, fixture.lambda, options);
    require(original.has_estimates() && scaled.has_estimates(),
            "scale-invariance fixture returned no estimates");
    require(original.status == scaled.status,
            "phenotype scaling changed the fit status");
    require_near(scaled.sigma_e2, original.sigma_e2 * scale * scale, 2e-10,
                 "sigma_e2 was not restored to phenotype units");
    require_near(scaled.sigma_g2, original.sigma_g2 * scale * scale, 2e-10,
                 "sigma_g2 was not restored to phenotype units");
    require_near(scaled.h2, original.h2, 2e-11,
                 "phenotype scaling changed h2");
    for (std::size_t i = 0; i < original.beta.size(); ++i) {
        require_near(scaled.beta[i], original.beta[i] * scale, 2e-10,
                     "beta was not restored to phenotype units");
    }
    for (std::size_t i = 0;
         i < original.beta_covariance_packed_lower.size(); ++i) {
        require_near(scaled.beta_covariance_packed_lower[i],
                     original.beta_covariance_packed_lower[i] * scale * scale,
                     3e-10,
                     "beta covariance was not restored to phenotype units");
    }
    const double expected_log_likelihood_shift =
        -static_cast<double>(fixture.y.size() - fixture.x.cols()) *
        std::log(scale);
    require_near(scaled.log_likelihood,
                 original.log_likelihood + expected_log_likelihood_shift,
                 2e-11,
                 "REML likelihood was not restored to phenotype units");
}

void test_residual_only_kkt_boundary() {
    constexpr std::size_t n = 14;
    spectra::reml::ColMajorMatrix x(n, 1);
    std::vector<double> lambda(n);
    std::vector<double> y(n, 1.5);
    for (std::size_t i = 0; i < n; ++i) {
        x(i, 0) = 1.0;
        lambda[i] = 0.1 * static_cast<double>(i + 1);
    }
    // The residual energy is confined to the two smallest GRM eigenvalues,
    // making the one-sided score for v_g at zero strictly negative.
    y[0] += 1.0;
    y[1] -= 1.0;
    const double expected_sigma_e2 = 2.0 / static_cast<double>(n - 1);
    const auto boundary = spectra::reml::evaluate_reml_spectral(
        y, x, lambda, std::sqrt(expected_sigma_e2), 0.0, false, true);
    require(boundary.valid, "residual-only boundary evaluation failed: " +
                                boundary.error);
    require(boundary.variance_score_g < -1e-3,
            "fixture does not have a negative one-sided genetic score");

    spectra::reml::RemlOptions options;
    options.max_iterations = 100;
    options.gradient_absolute_tolerance = 1e-10;
    options.gradient_relative_tolerance = 1e-11;
    const auto fitted =
        spectra::reml::fit_ai_reml_spectral(y, x, lambda, options);
    require(fitted.has_estimates(),
            "boundary AI-REML returned no estimates: " + fitted.error);
    require(fitted.status == spectra::reml::FitStatus::converged_boundary,
            std::string("residual-only optimum was not accepted as a KKT boundary: ") +
                spectra::reml::to_string(fitted.status) + " " + fitted.error);
    require(fitted.iterations <= 10,
            "near-boundary KKT check did not terminate promptly");
    require_near(fitted.sigma_g2, 0.0, 1e-15,
                 "genetic boundary must report exactly zero variance");
    require_near(fitted.h2, 0.0, 1e-15,
                 "genetic boundary must report exactly zero h2");
    require_near(fitted.sigma_e2, expected_sigma_e2, 2e-12,
                 "residual-only variance is not RSS/(n-p)");
    require_near(fitted.log_likelihood, boundary.log_likelihood, 2e-12,
                 "reported boundary likelihood is inconsistent");

    const double expected_se = std::sqrt(expected_sigma_e2 / n);
    const double expected_t = 1.5 / expected_se;
    const auto satterthwaite = spectra::reml::infer_fixed_effects_spectral(
        y, x, lambda, fitted,
        spectra::reml::FixedEffectTestMethod::satterthwaite);
    require(satterthwaite.status ==
                spectra::reml::FixedEffectInferenceStatus::boundary_ols,
            "Satterthwaite request at the boundary did not fall back to OLS");
    require(satterthwaite.method ==
                spectra::reml::FixedEffectTestMethod::ordinary_least_squares,
            "boundary inference did not report ordinary least squares");
    require(satterthwaite.coefficient_tests.size() == 1,
            "Satterthwaite coefficient test is missing");
    const auto& satt = satterthwaite.coefficient_tests.front();
    require(satt.valid, "Satterthwaite boundary test is invalid");
    require_near(satt.denominator_df, static_cast<double>(n - 1), 2e-11,
                 "Satterthwaite did not recover OLS residual df");
    require_near(satt.standard_error, expected_se, 2e-12,
                 "Satterthwaite did not recover OLS standard error");
    require_near(satt.statistic, expected_t, 2e-11,
                 "Satterthwaite did not recover OLS t statistic");

    const auto kr = spectra::reml::infer_fixed_effects_spectral(
        y, x, lambda, fitted,
        spectra::reml::FixedEffectTestMethod::kenward_roger);
    require(kr.status ==
                spectra::reml::FixedEffectInferenceStatus::boundary_ols,
            "KR request at the boundary did not fall back to OLS");
    require(kr.method ==
                spectra::reml::FixedEffectTestMethod::ordinary_least_squares,
            "boundary KR request did not report ordinary least squares");
    require(kr.coefficient_tests.size() == 1,
            "KR coefficient test is missing");
    const auto& kr_test = kr.coefficient_tests.front();
    require(kr_test.valid, "KR boundary test is invalid: " + kr.error);
    require_near(kr_test.denominator_df, static_cast<double>(n - 1), 2e-10,
                 "KR did not recover OLS residual df");
    require_near(kr_test.standard_error, expected_se, 2e-11,
                 "KR did not recover OLS standard error");
    require_near(kr_test.statistic, expected_t, 2e-10,
                 "boundary fallback did not recover OLS t statistic");
}

}  // namespace

int main() {
    try {
        test_score_matches_finite_difference();
        test_signed_parameterization_and_row_permutation();
        test_ai_reml_fit_improves_likelihood();
        test_probability_distributions_against_r();
        test_fixed_effect_inference_against_dense_reference();
        test_internal_phenotype_scaling_restores_original_units();
        test_residual_only_kkt_boundary();
        std::cout << "test_reml_synthetic: PASS\n";
        return EXIT_SUCCESS;
    } catch (const std::exception& exception) {
        std::cerr << "test_reml_synthetic: FAIL: " << exception.what() << '\n';
        return EXIT_FAILURE;
    }
}