#include "spectra_reml/grm.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#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 tolerance,
                  const std::string& message) {
    if (std::abs(actual - expected) > tolerance) {
        throw std::runtime_error(message + ": actual=" +
                                 std::to_string(actual) +
                                 ", expected=" + std::to_string(expected));
    }
}

class TemporaryDirectory {
public:
    TemporaryDirectory() {
        const auto seed =
            std::chrono::high_resolution_clock::now().time_since_epoch().count();
        path_ = std::filesystem::temp_directory_path() /
                ("spectra_reml_grm_test_" + std::to_string(seed));
        std::filesystem::create_directories(path_);
    }
    ~TemporaryDirectory() {
        std::error_code ignored;
        std::filesystem::remove_all(path_, ignored);
    }
    const std::filesystem::path& path() const noexcept { return path_; }

private:
    std::filesystem::path path_;
};

void write_floats(const std::filesystem::path& path,
                  const std::vector<float>& values) {
    std::ofstream output(path, std::ios::binary | std::ios::trunc);
    output.write(reinterpret_cast<const char*>(values.data()),
                 static_cast<std::streamsize>(values.size() * sizeof(float)));
    if (!output) {
        throw std::runtime_error("failed to create GRM fixture");
    }
}

void test_read_and_decompose() {
    TemporaryDirectory temporary;
    const auto grm_path = temporary.path() / "fixture.grm.bin";
    const auto id_path = temporary.path() / "fixture.grm.id";
    // Symmetric positive-definite matrix, packed as GCTA lower triangle.
    const std::vector<float> packed = {1.2F, 0.2F, 1.1F,
                                       -0.1F, 0.15F, 0.9F};
    write_floats(grm_path, packed);
    {
        std::ofstream ids(id_path);
        ids << "F1 I1\nF2 I2\nF3 I3\n";
    }
    spectra::reml::validate_grm_id_count(id_path, 3);

    const auto matrix = spectra::reml::read_gcta_grm_lower_triangle(grm_path, 3);
    require_near(matrix(0, 0), 1.2, 1e-7, "wrong GRM diagonal");
    require_near(matrix(1, 0), 0.2, 1e-7, "wrong GRM lower triangle");
    require_near(matrix(0, 1), 0.2, 1e-7, "GRM was not symmetrized");
    require_near(matrix(2, 0), -0.1, 1e-7, "wrong GRM packed order");
    require_near(matrix(1, 2), 0.15, 1e-7, "wrong GRM upper triangle");

    const auto spectral =
        spectra::reml::read_and_diagonalize_gcta_grm(grm_path, 3);
    require(spectral.eigenvalues.size() == 3, "wrong eigenvalue count");
    require(std::is_sorted(spectral.eigenvalues.begin(),
                           spectral.eigenvalues.end()),
            "LAPACK eigenvalues are not ascending");
    require(spectral.minimum_eigenvalue > 0.0,
            "positive-definite fixture has non-positive eigenvalue");

    // Reconstruct U diag(lambda) U' and compare to the original matrix.
    for (std::size_t row = 0; row < 3; ++row) {
        for (std::size_t col = 0; col < 3; ++col) {
            double reconstructed = 0.0;
            for (std::size_t axis = 0; axis < 3; ++axis) {
                reconstructed += spectral.eigenvectors(row, axis) *
                                 spectral.eigenvalues[axis] *
                                 spectral.eigenvectors(col, axis);
            }
            require_near(reconstructed, matrix(row, col), 2e-12,
                         "eigendecomposition does not reconstruct GRM");
        }
    }
}

void test_wrong_file_size_is_rejected() {
    TemporaryDirectory temporary;
    const auto grm_path = temporary.path() / "short.grm.bin";
    write_floats(grm_path, {1.0F, 0.0F});
    bool rejected = false;
    try {
        (void)spectra::reml::read_gcta_grm_lower_triangle(grm_path, 2);
    } catch (const std::runtime_error&) {
        rejected = true;
    }
    require(rejected, "short GRM file was not rejected");
}

}  // namespace

int main() {
    try {
        test_read_and_decompose();
        test_wrong_file_size_is_rejected();
        std::cout << "test_grm_io: PASS\n";
        return EXIT_SUCCESS;
    } catch (const std::exception& exception) {
        std::cerr << "test_grm_io: FAIL: " << exception.what() << '\n';
        return EXIT_FAILURE;
    }
}