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

#include <catch2/catch_all.hpp>
#include "../common/GridScanSettings.h"
#include "../common/JFJochException.h"

TEST_CASE("GridScanSettings basic construction", "[GridScanSettings]") {
    SECTION("Valid parameters") {
        REQUIRE_NOTHROW(GridScanSettings(10, 1.0f, 2.0f, false, false));
        REQUIRE_NOTHROW(GridScanSettings(10, 1.0f, -2.0f, true, false));
        REQUIRE_NOTHROW(GridScanSettings(10, -1.0f, 2.0f, false, true));
        REQUIRE_NOTHROW(GridScanSettings(10, 1.0f, 2.0f, true, true));
    }

    SECTION("Invalid parameters") {
// Negative n_fast
        REQUIRE_THROWS_AS(GridScanSettings(-5, 1.0f, 2.0f, false, false), JFJochException);

// Zero n_fast
        REQUIRE_THROWS_AS(GridScanSettings(0, 1.0f, 2.0f, false, false), JFJochException);

// Zero grid_step_fast_um
        REQUIRE_THROWS_AS(GridScanSettings(10, 0.0f, 2.0f, false, false), JFJochException);

// Zero grid_step_slow_um
        REQUIRE_THROWS_AS(GridScanSettings(10, 1.0f, 0.0f, false, false), JFJochException);
    }
}

TEST_CASE("GridScanSettings image number", "[GridScanSettings]") {
    GridScanSettings grid(5, 1.5f, 2.5f, false, false);
    CHECK_THROWS(grid.ImageNum(-1));

    grid.ImageNum(0);
    CHECK(grid.GetNSlow() == 0);
    CHECK(grid.GetNFast() == 5);

    grid.ImageNum(1);
    CHECK(grid.GetNSlow() == 1);
    CHECK(grid.GetNFast() == 5);

    grid.ImageNum(4);
    CHECK(grid.GetNSlow() == 1);
    CHECK(grid.GetNFast() == 5);

    grid.ImageNum(6);
    CHECK(grid.GetNSlow() == 2);
    CHECK(grid.GetNFast() == 5);

    grid.ImageNum(13);
    CHECK(grid.GetNSlow() == 3);
    CHECK(grid.GetNFast() == 5);
}

TEST_CASE("GridScanSettings horizontal scan", "[GridScanSettings]") {
// Create a grid with 5 elements in fast direction, 1.5 um step in fast, 2.5 um in slow
    GridScanSettings grid(5, 1.5f, 2.5f, false, false);

    SECTION("Grid sizes in steps") {
        REQUIRE(grid.GetGridSizeX_step() == 5);
        REQUIRE(grid.GetGridSizeY_step() == 1);
        REQUIRE(grid.GetGridSizeX_um() == Catch::Approx(5 * 1.5f));
        REQUIRE(grid.GetGridSizeY_um() == Catch::Approx(1 * 2.5f));
    }

    grid.ImageNum(40);
    SECTION("Grid sizes in steps") {
        REQUIRE(grid.GetGridSizeX_step() == 5);
        REQUIRE(grid.GetGridSizeY_step() == 8);
        REQUIRE(grid.GetGridSizeX_um() == Catch::Approx(5 * 1.5f));
        REQUIRE(grid.GetGridSizeY_um() == Catch::Approx(8 * 2.5f));
    }

    SECTION("Element positions in standard raster") {
// First row
        REQUIRE(grid.GetElementPosX_step(0) == 0);
        REQUIRE(grid.GetElementPosY_step(0) == 0);
        REQUIRE(grid.GetElementPosX_step(1) == 1);
        REQUIRE(grid.GetElementPosY_step(1) == 0);
        REQUIRE(grid.GetElementPosX_step(4) == 4);
        REQUIRE(grid.GetElementPosY_step(4) == 0);

// Second row
        REQUIRE(grid.GetElementPosX_step(5) == 0);
        REQUIRE(grid.GetElementPosY_step(5) == 1);
        REQUIRE(grid.GetElementPosX_step(9) == 4);
        REQUIRE(grid.GetElementPosY_step(9) == 1);

// Positions in um
        REQUIRE(grid.GetElementPosX_um(4) == Catch::Approx(6.0f));  // 4 * 1.5
        REQUIRE(grid.GetElementPosY_um(9) == Catch::Approx(2.5f));  // 1 * 2.5
    }
}

TEST_CASE("GridScanSettings vertical scan", "[GridScanSettings]") {
// Create a vertical grid with 5 elements in fast direction
    GridScanSettings grid(5, 2.5f, 1.5f, false, true);
    grid.ImageNum(15);
    SECTION("Grid sizes in steps") {
// For vertical scan, fast = Y
        REQUIRE(grid.GetGridSizeY_step() == 5);
        REQUIRE(grid.GetGridSizeX_step() == 3);
    }

    SECTION("Grid sizes in um") {
        REQUIRE(grid.GetGridSizeY_um() == Catch::Approx(5 * 1.5f));
        REQUIRE(grid.GetGridSizeX_um() == Catch::Approx(3 * 2.5f));
    }

    SECTION("Grid steps") {
        REQUIRE(grid.GetGridStepY_um() == Catch::Approx(1.5f));
        REQUIRE(grid.GetGridStepX_um() == Catch::Approx(2.5f));
    }

    SECTION("Element positions in vertical scan") {
// First column
        REQUIRE(grid.GetElementPosX_step(0) == 0);
        REQUIRE(grid.GetElementPosY_step(0) == 0);
        REQUIRE(grid.GetElementPosX_step(1) == 0);
        REQUIRE(grid.GetElementPosY_step(1) == 1);
        REQUIRE(grid.GetElementPosX_step(4) == 0);
        REQUIRE(grid.GetElementPosY_step(4) == 4);

// Second column
        REQUIRE(grid.GetElementPosX_step(5) == 1);
        REQUIRE(grid.GetElementPosY_step(5) == 0);
        REQUIRE(grid.GetElementPosX_step(9) == 1);
        REQUIRE(grid.GetElementPosY_step(9) == 4);
    }
}

TEST_CASE("GridScanSettings snake raster scan", "[GridScanSettings]") {
// Create a grid with snake pattern - horizontal
    GridScanSettings grid(5, 1.5f, 2.5f, true, false);
    grid.ImageNum(50);

    SECTION("Element positions in snake raster") {
// First row (left to right)
        REQUIRE(grid.GetElementPosX_step(0) == 0);
        REQUIRE(grid.GetElementPosY_step(0) == 0);
        REQUIRE(grid.GetElementPosX_step(4) == 4);
        REQUIRE(grid.GetElementPosY_step(4) == 0);

// Second row (right to left)
        REQUIRE(grid.GetElementPosX_step(5) == 4);
        REQUIRE(grid.GetElementPosY_step(5) == 1);
        REQUIRE(grid.GetElementPosX_step(9) == 0);
        REQUIRE(grid.GetElementPosY_step(9) == 1);

// Third row (left to right)
        REQUIRE(grid.GetElementPosX_step(10) == 0);
        REQUIRE(grid.GetElementPosY_step(10) == 2);
        REQUIRE(grid.GetElementPosX_step(14) == 4);
        REQUIRE(grid.GetElementPosY_step(14) == 2);
    }
}

TEST_CASE("GridScanSettings vertical snake raster scan", "[GridScanSettings]") {
// Create a grid with snake pattern - vertical
    GridScanSettings grid(5, 1.5f, 2.5f, true, true);
    grid.ImageNum(50);

    SECTION("Element positions in vertical snake raster") {
// First column (top to bottom)
        REQUIRE(grid.GetElementPosX_step(0) == 0);
        REQUIRE(grid.GetElementPosY_step(0) == 0);
        REQUIRE(grid.GetElementPosX_step(4) == 0);
        REQUIRE(grid.GetElementPosY_step(4) == 4);

// Second column (bottom to top)
        REQUIRE(grid.GetElementPosX_step(5) == 1);
        REQUIRE(grid.GetElementPosY_step(5) == 4);
        REQUIRE(grid.GetElementPosX_step(9) == 1);
        REQUIRE(grid.GetElementPosY_step(9) == 0);

// Third column (top to bottom)
        REQUIRE(grid.GetElementPosX_step(10) == 2);
        REQUIRE(grid.GetElementPosY_step(10) == 0);
        REQUIRE(grid.GetElementPosX_step(14) == 2);
        REQUIRE(grid.GetElementPosY_step(14) == 4);
    }
}

TEST_CASE("GridScanSettings negative fast step", "[GridScanSettings]") {
// Create a grid with negative step sizes
    GridScanSettings grid(5, -1.5f, 2.5f, false, false);
    grid.ImageNum(35);

    SECTION("Element positions") {
// First column
        REQUIRE(grid.GetElementPosX_step(0) == 4);
        REQUIRE(grid.GetElementPosY_step(0) == 0);
        REQUIRE(grid.GetElementPosX_step(4) == 0);
        REQUIRE(grid.GetElementPosY_step(4) == 0);

// Second column
        REQUIRE(grid.GetElementPosX_step(5) == 4);
        REQUIRE(grid.GetElementPosY_step(5) == 1);
        REQUIRE(grid.GetElementPosX_step(9) == 0);
        REQUIRE(grid.GetElementPosY_step(9) == 1);
    }

    SECTION("Grid sizes with negative steps") {
// Grid sizes should use absolute values
        REQUIRE(grid.GetGridSizeX_um() == Catch::Approx(5 * 1.5f));
        REQUIRE(grid.GetGridSizeY_um() == Catch::Approx(7 * 2.5f));
    }

    SECTION("Element positions with negative steps") {
        REQUIRE(grid.GetElementPosX_um(0) == Catch::Approx(4 * 1.5f));
        REQUIRE(grid.GetElementPosX_um(4) == Catch::Approx(0));
        REQUIRE(grid.GetElementPosY_um(5) == Catch::Approx(2.5f));
    }

    SECTION("Grid steps maintain sign") {
        REQUIRE(grid.GetGridStepX_um() == Catch::Approx(-1.5f));
        REQUIRE(grid.GetGridStepY_um() == Catch::Approx(2.5f));
    }
}

TEST_CASE("GridScanSettings negative slow step", "[GridScanSettings]") {
// Create a grid with negative step sizes
    GridScanSettings grid(5, 1.5f, -2.5f, false, false);
    grid.ImageNum(35);

    SECTION("Element positions") {
// First column
        REQUIRE(grid.GetElementPosX_step(0) == 0);
        REQUIRE(grid.GetElementPosY_step(0) == 6);
        REQUIRE(grid.GetElementPosX_step(4) == 4);
        REQUIRE(grid.GetElementPosY_step(4) == 6);

// Second column
        REQUIRE(grid.GetElementPosX_step(5) == 0);
        REQUIRE(grid.GetElementPosY_step(5) == 5);
        REQUIRE(grid.GetElementPosX_step(9) == 4);
        REQUIRE(grid.GetElementPosY_step(9) == 5);
    }

    SECTION("Grid sizes with negative steps") {
// Grid sizes should use absolute values
        REQUIRE(grid.GetGridSizeX_um() == Catch::Approx(5 * 1.5f));
        REQUIRE(grid.GetGridSizeY_um() == Catch::Approx(7 * 2.5f));
    }

    SECTION("Element positions with negative steps") {
// Positions should maintain sign
        REQUIRE(grid.GetElementPosX_um(4) == Catch::Approx(4 * 1.5f));
        REQUIRE(grid.GetElementPosY_um(5) == Catch::Approx(5 * 2.5f));
    }

    SECTION("Grid steps maintain sign") {
        REQUIRE(grid.GetGridStepX_um() == Catch::Approx(1.5f));
        REQUIRE(grid.GetGridStepY_um() == Catch::Approx(-2.5f));
    }
}

TEST_CASE("GridScanSettings::Rearrange functionality tests", "[gridscan][rearrange]") {
    SECTION("Basic rearrangement in standard grid configuration") {
        // Create a basic 3x3 grid with standard parameters
        GridScanSettings grid(3, 1.0f, 1.0f, false, false);
        grid.ImageNum(9); // 3x3 grid

        // Create test input data [0,1,2,3,4,5,6,7,8]
        std::vector<float> input(9);
        for (int i = 0; i < 9; i++) {
            input[i] = static_cast<float>(i);
        }

        // Rearrange the data
        std::vector<float> output = grid.Rearrange(input);

        // In standard mode (no snake, no vertical), data should be mapped correctly
        // Each row in source (0,1,2), (3,4,5), (6,7,8)
        // should appear in the same positions in the output
        REQUIRE(output.size() == 9);
        for (int i = 0; i < 9; i++) {
            REQUIRE(output[i] == Catch::Approx(input[i]));
        }
    }

    SECTION("Snake raster scan rearrangement") {
        // Create a grid with snake pattern (alternating row directions)
        GridScanSettings grid(3, 1.0f, 1.0f, true, false);
        grid.ImageNum(9); // 3x3 grid

        // Create test input data [0,1,2,3,4,5,6,7,8]
        std::vector<float> input(9);
        for (int i = 0; i < 9; i++) {
            input[i] = static_cast<float>(i);
        }

        // Rearrange the data
        std::vector<float> output = grid.Rearrange(input);

        // In snake mode, odd rows should be reversed
        // Input: [0,1,2,3,4,5,6,7,8]
        // Expected in rows: [0,1,2], [5,4,3], [6,7,8]
        // As flattened array: [0,1,2,5,4,3,6,7,8]
        REQUIRE(output.size() == 9);
        std::vector<float> expected = {0, 1, 2, 5, 4, 3, 6, 7, 8};
        for (size_t i = 0; i < expected.size(); i++) {
            REQUIRE(output[i] == Catch::Approx(expected[i]));
        }
    }

    SECTION("Vertical scan rearrangement") {
        // Test vertical scan mode
        GridScanSettings grid(3, 1.0f, 1.0f, false, true);
        grid.ImageNum(9); // 3x3 grid

        // Create test input data [0,1,2,3,4,5,6,7,8]
        std::vector<float> input(9);
        for (int i = 0; i < 9; i++) {
            input[i] = static_cast<float>(i);
        }

        // Rearrange the data
        std::vector<float> output = grid.Rearrange(input);

        // In vertical mode, scan goes down columns first
        // Input: [0,1,2,3,4,5,6,7,8]
        // Expected as columns: [0,3,6], [1,4,7], [2,5,8]
        // But indexed by row-major order, so positions transform
        REQUIRE(output.size() == 9);

        // Verify transformations by checking specific indices
        // The exact values depend on how vertical scan is implemented
        // These assertions may need adjustment based on implementation
        REQUIRE(output[0] == Catch::Approx(input[0]));
        REQUIRE(output[3] == Catch::Approx(input[1]));
        REQUIRE(output[6] == Catch::Approx(input[2]));
        REQUIRE(output[1] == Catch::Approx(input[3]));
    }

    SECTION("Mixed configuration - snake and vertical scan") {
        // Test combination of snake and vertical
        GridScanSettings grid(3, 1.0f, 1.0f, true, true);
        grid.ImageNum(9); // 3x3 grid

        // Create test input data [0,1,2,3,4,5,6,7,8]
        std::vector<float> input(9);
        for (int i = 0; i < 9; i++) {
            input[i] = static_cast<float>(i);
        }

        // Rearrange the data
        std::vector<float> output = grid.Rearrange(input);
        // 0 5 6
        // 1 4 7
        // 2 3 8
        std::vector<float> expected = {0, 5, 6, 1, 4, 7, 2, 3, 8};
        // Complex case with both features enabled
        REQUIRE(output.size() == 9);
        // Specific values would depend on implementation but ensure we get a valid output
        for (size_t i = 0; i < expected.size(); i++) {
            CHECK(output[i] == Catch::Approx(expected[i]));
        }
    }

    SECTION("Negative grid steps") {
        // Test with negative grid steps
        GridScanSettings grid(3, -1.0f, -1.0f, false, false);
        grid.ImageNum(9); // 3x3 grid

        // Create test input data [0,1,2,3,4,5,6,7,8]
        std::vector<float> input(9);
        for (int i = 0; i < 9; i++) {
            input[i] = static_cast<float>(i);
        }

        // Rearrange the data
        std::vector<float> output = grid.Rearrange(input);

        // Both fast and slow directions are negative, so mapping inverts both ways
        // Expected: [8,7,6,5,4,3,2,1,0]
        REQUIRE(output.size() == 9);
        for (int i = 0; i < 9; i++) {
            REQUIRE(output[i] == Catch::Approx(input[8-i]));
        }
    }

    SECTION("Input smaller than grid") {
        // Test case where input is smaller than grid
        GridScanSettings grid(3, 1.0f, 1.0f, false, false);
        grid.ImageNum(9); // 3x3 grid

        // Create test input data [0,1,2,3,4] (shorter than grid)
        std::vector<float> input = {0, 1, 2, 3, 4};

        // Rearrange the data
        std::vector<float> output = grid.Rearrange(input, 172.0f);

        // Output should be full size with NANs for missing data
        REQUIRE(output.size() == 9);
        for (int i = 0; i < 5; i++) {
            REQUIRE(output[i] == Catch::Approx(input[i]));
        }
        for (int i = 5; i < 9; i++) {
            REQUIRE(output[i] == 172.0f);
        }
    }

    SECTION("Input larger than grid") {
        // Test case where input is larger than grid
        GridScanSettings grid(2, 1.0f, 1.0f, false, false);
        grid.ImageNum(4); // 2x2 grid

        // Create test input with more elements than grid
        std::vector<float> input = {0, 1, 2, 3, 4, 5};

        // Rearrange the data
        std::vector<float> output = grid.Rearrange(input);

        // Output should match grid size and ignore extra input
        REQUIRE(output.size() == 4);
        for (int i = 0; i < 4; i++) {
            REQUIRE(output[i] == Catch::Approx(input[i]));
        }
    }
}