SpectraREML/include/spectra_reml/batch_io.hpp

#pragma once

#include "spectra_reml/grm.hpp"
#include "spectra_reml/types.hpp"

#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <string>
#include <vector>

namespace spectra::reml {

struct BatchInputPaths {
    std::filesystem::path grm_bin;
    std::filesystem::path grm_id;
    std::filesystem::path base_x_f64;
    std::filesystem::path phenotype_f64;
    std::filesystem::path extra_covariate_f32;
    std::filesystem::path tasks_tsv;
    std::filesystem::path extra_offsets_i64;
    std::filesystem::path extra_indices_i32;
    std::filesystem::path contrast_matrix_f64;
    std::filesystem::path contrast_metadata_tsv;
    std::filesystem::path output_directory;
};

struct BatchDimensions {
    std::size_t sample_count = 0;
    std::size_t base_covariate_count = 0;
    std::size_t phenotype_row_count = 0;
    std::size_t extra_covariate_row_count = 0;
    std::size_t contrast_count = 0;
};

struct BatchOptions {
    std::size_t block_size = 256;
    std::size_t outer_threads = 1;
    bool resume = false;
    bool overwrite = false;
    RemlOptions reml;
};

[[nodiscard]] ColMajorMatrix read_row_major_f64_matrix(
    const std::filesystem::path& path, std::size_t rows, std::size_t cols);

[[nodiscard]] std::vector<RemlTask> read_tasks(
    const std::filesystem::path& tasks_tsv,
    const std::filesystem::path& extra_offsets_i64,
    const std::filesystem::path& extra_indices_i32,
    const BatchDimensions& dimensions);

// Runs one complete task set.  The GRM and common variables are transformed
// once; phenotypes are transformed in blocks and tasks are fitted in parallel.
// Each completed block contains:
//   block_NNNNNN.summary.tsv
//   block_NNNNNN.beta.f64.bin
//   block_NNNNNN.cov.f64.bin
//   block_NNNNNN.complete
// The .complete marker is renamed last and is the sole resume criterion.
void run_task_batch(const BatchInputPaths& paths,
                    const BatchDimensions& dimensions,
                    const BatchOptions& options);

}  // namespace spectra::reml