// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <cmath>
#include "GridScanSettings.h"
#include "JFJochException.h"

GridScanSettings::GridScanSettings(int64_t in_n_fast,
                                   float in_grid_step_x_um,
                                   float in_grid_step_y_um,
                                   bool in_snake_raster_scan,
                                   bool in_vertical_scan)
        : n_fast(in_n_fast),
          n_slow(1),
          n_elem(in_n_fast),
          grid_elem_x_um(in_grid_step_x_um),
          grid_elem_y_um(in_grid_step_y_um),
          snake_scan(in_snake_raster_scan),
          vertical_scan(in_vertical_scan) {
    if (n_fast <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Number of elements in fast direction must be positive");
    if (grid_elem_x_um == 0.0f || grid_elem_y_um == 0.0f)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Grid elements must be non-zero");
}

GridScanSettings &GridScanSettings::ImageNum(int64_t input) {
    if (input < 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Number of elements in slow direction must be positive");

    n_slow = (input + n_fast - 1) / n_fast;
    n_elem = n_slow * n_fast;
    return *this;
}

int64_t GridScanSettings::GetGridSizeX_step() const {
    return vertical_scan ? n_slow : n_fast;
}

int64_t GridScanSettings::GetGridSizeY_step() const {
    return vertical_scan ? n_fast : n_slow;
}

float GridScanSettings::GetGridSizeX_um() const {
    return fabsf(GetGridStepX_um()) * static_cast<float>(GetGridSizeX_step());
}

float GridScanSettings::GetGridSizeY_um() const {
    return fabsf(GetGridStepY_um()) * static_cast<float>(GetGridSizeY_step());
}

int64_t GridScanSettings::GetElementPosX_step(int64_t elem_number) const {
    if (vertical_scan)
        return GetElementPosSlow_step(elem_number);
    else
        return GetElementPosFast_step(elem_number);
}

int64_t GridScanSettings::GetElementPosY_step(int64_t elem_number) const {
    if (vertical_scan)
        return GetElementPosFast_step(elem_number);
    else
        return GetElementPosSlow_step(elem_number);
}

float GridScanSettings::GetElementPosX_um(int64_t elem_number) const {
    return static_cast<float>(GetElementPosX_step(elem_number)) * fabs(GetGridStepX_um());
}

float GridScanSettings::GetElementPosY_um(int64_t elem_number) const {
    return static_cast<float>(GetElementPosY_step(elem_number)) * fabs(GetGridStepY_um());
}

float GridScanSettings::GetGridStepX_um() const {
    return grid_elem_x_um;
}

float GridScanSettings::GetGridStepY_um() const {
    return grid_elem_y_um;
}

int64_t GridScanSettings::GetElementPosFast_step(int64_t image_number) const {
    if ((image_number < 0) || (image_number >= n_elem))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Image out of bounds");

    int64_t slow = GetElementPosSlow_step(image_number);
    int64_t fast = image_number % n_fast;
    if (GetGridElemFast_um() < 0)
        fast = (n_fast - 1) - fast;
    if (snake_scan && (slow % 2 == 1))
        fast = (n_fast - 1) - fast;
    return fast;
}

int64_t GridScanSettings::GetElementPosSlow_step(int64_t image_number) const {
    if ((image_number < 0) || (image_number >= n_elem))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Image out of bounds");

    int64_t slow = image_number / n_fast;
    if (GetGridElemSlow_um() < 0)
        slow = (n_slow - 1) - slow;
    return slow;
}

std::vector<double> GridScanSettings::GetXContainer_m(int64_t max_image_number) const {
    std::vector<double> pos_container(max_image_number);
    for (int32_t i = 0; i < max_image_number; i++)
        pos_container[i] = GetElementPosX_um(i) * 1e-6;
    return pos_container;
}

std::vector<double> GridScanSettings::GetYContainer_m(int64_t max_image_number) const {
    std::vector<double> pos_container(max_image_number);
    for (int32_t i = 0; i < max_image_number; i++)
        pos_container[i] = GetElementPosY_um(i) * 1e-6;
    return pos_container;
}

int64_t GridScanSettings::GetNFast() const {
    return n_fast;
}

int64_t GridScanSettings::GetNSlow() const {
    return n_slow;
}

int64_t GridScanSettings::GetNElem() const {
    return n_elem;
}

float GridScanSettings::GetGridElemFast_um() const {
    if (vertical_scan)
        return grid_elem_y_um;
    else
        return grid_elem_x_um;
}

float GridScanSettings::GetGridElemSlow_um() const {
    if (vertical_scan)
        return grid_elem_x_um;
    return grid_elem_y_um;
}

bool GridScanSettings::IsSnakeScan() const {
    return snake_scan;
}

bool GridScanSettings::IsVerticalScan() const {
    return vertical_scan;
}

std::vector<float> GridScanSettings::Rearrange(const std::vector<float> &input, float fill_value) const {
    std::vector<float> output(n_elem, fill_value);
    for (int64_t i = 0; i < std::min<int64_t>(input.size(), n_elem); i++)
        output.at(Rearrange(i)) = input[i];

    return output;
}

std::vector<int64_t> GridScanSettings::Rearrange(const std::vector<int64_t> &input, uint64_t fill_value) const {
    std::vector<int64_t> output(n_elem, fill_value);
    for (int64_t i = 0; i < std::min<int64_t>(input.size(), n_elem); i++)
        output.at(Rearrange(i)) = input[i];
    return output;
}

int64_t GridScanSettings::Rearrange(int64_t image_number) const {
    return GetElementPosY_step(image_number) * GetGridSizeX_step()
            + GetElementPosX_step(image_number);
}