#include "testlib.h"
#include <queue>
#include <ranges>
#include <unordered_set>

enum Color {
    RED, BLUE
};

std::ostream& operator<<(std::ostream& out, const Color color) {
    return out << (color == RED ? 'R' : 'B');
}

struct Cell {
    uint32_t coordinate;
    uint32_t number_of_chips;
    Color color;

    bool operator<(const Cell& other) const {
        return coordinate < other.coordinate;
    }
};

std::ostream& operator<<(std::ostream& out, const Cell cell) {
    return out << cell.coordinate << ' ' << cell.number_of_chips << ' ' << cell.color;
}

struct Test {
    std::vector<Cell> cells;
};

std::ostream& operator<<(std::ostream& out, const Test& test) {
    out << test.cells.size() << '\n';
    std::ranges::copy(test.cells, std::ostream_iterator<Cell>(out, "\n"));
    return out;
}

struct MultiTest {
    std::vector<Test> tests;
};

std::ostream& operator<<(std::ostream& out, const MultiTest& multitest) {
    out << multitest.tests.size() << '\n';
    std::ranges::copy(multitest.tests, std::ostream_iterator<Test>(out, ""));
    return out;
}


std::vector<uint32_t> get_order(const uint32_t upper_coordinate_bound, const uint32_t upper_n_bound, const std::string& order_type) {
    if (order_type == "increasing")
        return std::ranges::views::iota(static_cast<uint32_t>(1), upper_coordinate_bound + 1) | std::ranges::to<std::vector<uint32_t>>();
    else if (order_type == "decreasing")
        return std::ranges::views::iota(static_cast<uint32_t>(1), upper_coordinate_bound + 1) | std::ranges::views::reverse | std::ranges::to<std::vector<uint32_t>>();
    else if (order_type == "jumping") {
        std::vector<uint32_t> v = {1, upper_coordinate_bound};
        std::queue<std::pair<uint32_t, uint32_t>> segments;
        if (upper_coordinate_bound >= 3)
            segments.emplace(2, upper_coordinate_bound);
        for (uint32_t segment_index = 2; segment_index < upper_n_bound && !segments.empty(); ++segment_index) {
            std::pair<uint32_t, uint32_t> segment = segments.front();
            segments.pop();
            const uint32_t m = (segment.first + segment.second) / 2;
            v.push_back(m);
            if (segment.first < m)
                segments.emplace(segment.first, m);
            if (m + 1 < segment.second)
                segments.emplace(m + 1, segment.second);
        }
        return v;
    } else {
        ensuref(false, "%s", std::format(R"(Unknown order type: "{}")", order_type).c_str());
        return {};
    }
}

std::vector<uint32_t> decompose_into_sum(uint32_t s, uint32_t n) {
    std::vector<bool> balls_and_septums(s + n - 1);
    for (uint32_t septum_index = 0; septum_index < n - 1; ++septum_index)
        balls_and_septums[septum_index] = true;
    shuffle(balls_and_septums.begin(), balls_and_septums.end());
    std::vector<uint32_t> partition(n);
    uint32_t current_addend_index = 0;
    for (bool cur : balls_and_septums)
        if (cur)
            ++current_addend_index;
        else
            ++partition[current_addend_index];
    return partition;
}

MultiTest generate_test() {
    std::string type = opt<std::string>("type");
    MultiTest multitest{};
    if (type == "few_chips") {
        uint32_t upper_coordinate_bound = opt<uint32_t>("upper_coordinate_bound");
        uint32_t upper_total_n_bound = opt<uint32_t>("upper_total_n_bound");
        uint32_t upper_n_bound = opt<uint32_t>("upper_n_bound");
        uint32_t number_of_chips = opt<uint32_t>("number_of_chips");
        uint32_t lower_n_tests_bound = opt<uint32_t>("lower_n_tests_bound");
        uint32_t upper_n_tests_bound = opt<uint32_t>("upper_n_tests_bound");
        std::string order_type = opt<std::string>("order_type");
        std::vector<uint32_t> order = get_order(upper_coordinate_bound, upper_n_bound, order_type);
        uint32_t total_n = 0;
        std::vector<int> coordinates(number_of_chips);
        do {
            if (coordinates.front() != coordinates.back()) {
                std::map<uint32_t, uint32_t> cells_map;
                for (uint32_t coordinate : coordinates)
                    ++cells_map[coordinate];
                std::vector<std::pair<uint32_t, uint32_t>> cells = cells_map | std::ranges::to<std::vector<std::pair<uint32_t, uint32_t>>>();
                uint32_t M = std::min<uint32_t>(30, cells.size());
                for (uint32_t mask = 1; (mask + 1 < static_cast<uint32_t>(1) << M) && total_n < upper_total_n_bound && multitest.tests.size() < upper_n_tests_bound; ++mask) {
                    total_n += cells.size();
                    multitest.tests.emplace_back();
                    multitest.tests.back().cells.resize(cells.size());
                    for (int cell_index = 0; cell_index < cells.size(); ++cell_index) {
                        multitest.tests.back().cells[cell_index] = {
                                .coordinate = order[cells[cell_index].first],
                                .number_of_chips = cells[cell_index].second,
                                .color = cell_index < M && (mask >> cell_index & 1) ? BLUE : RED,
                        };
                    }
                    sort(multitest.tests.back().cells.begin(), multitest.tests.back().cells.end());
                }
            }
            uint32_t not_equal_to_previous = 0;
            for (uint32_t chip_index = 1; chip_index < number_of_chips; ++chip_index)
                if (coordinates[chip_index] != coordinates[chip_index - 1])
                    not_equal_to_previous = chip_index;
            ++coordinates[not_equal_to_previous];
            for (uint32_t chip_index = not_equal_to_previous + 1; chip_index < number_of_chips; ++chip_index)
                coordinates[chip_index] = 0;
        } while (coordinates.front() < upper_n_bound && total_n < upper_total_n_bound && multitest.tests.size() < upper_n_tests_bound);
        while (!(multitest.tests.empty() || total_n <= upper_total_n_bound && multitest.tests.size() <= upper_n_tests_bound))
            multitest.tests.pop_back();
        ensuref(multitest.tests.size() >= lower_n_tests_bound, "%s", std::format("failed to create a multitest of size >= {}, got {} instead",
                                                                                 lower_n_tests_bound, multitest.tests.size()).c_str());
        shuffle(multitest.tests.begin(), multitest.tests.end());
        return multitest;
    } else if (type == "random") {
        uint32_t lower_n_tests_bound = opt<uint32_t>("lower_n_tests_bound");
        uint32_t upper_n_tests_bound = opt<uint32_t>("upper_n_tests_bound");
        uint32_t lower_m_bound = opt<uint32_t>("lower_m_bound");
        uint32_t upper_m_bound = opt<uint32_t>("upper_m_bound");
        uint32_t lower_x_bound = opt<uint32_t>("lower_coordinate_bound");
        uint32_t upper_x_bound = opt<uint32_t>("upper_coordinate_bound");
        uint32_t lower_total_n_bound = opt<uint32_t>("lower_total_n_bound");
        uint32_t upper_total_n_bound = opt<uint32_t>("upper_total_n_bound");
        upper_n_tests_bound = std::min(upper_n_tests_bound, upper_total_n_bound / 2);
        uint32_t n_tests = rnd.next(lower_n_tests_bound, upper_n_tests_bound);
        lower_total_n_bound = std::max(lower_total_n_bound, 2 * n_tests);
        uint32_t total_n = rnd.next(lower_total_n_bound, upper_total_n_bound);
        std::vector<uint32_t> test_sizes = decompose_into_sum(total_n - 2 * n_tests, n_tests);
        for (uint32_t& test_size : test_sizes)
            test_size += 2;
        multitest.tests.resize(n_tests);
        for (uint32_t test_index = 0; test_index < n_tests; ++test_index) {
            multitest.tests[test_index].cells.resize(test_sizes[test_index]);
            uint32_t number_of_reds = rnd.next(1u, test_sizes[test_index] - 1);
            std::unordered_set<uint32_t> taken_coordinates;
            for (uint32_t cell_index = 0; cell_index < test_sizes[test_index]; ++cell_index) {
                multitest.tests[test_index].cells[cell_index].number_of_chips = rnd.next(lower_m_bound, upper_m_bound);
                multitest.tests[test_index].cells[cell_index].color = cell_index < number_of_reds ? RED : BLUE;
                do {
                    multitest.tests[test_index].cells[cell_index].coordinate = rnd.next(lower_x_bound, upper_x_bound);
                    taken_coordinates.insert(multitest.tests[test_index].cells[cell_index].coordinate);
                } while (taken_coordinates.size() <= cell_index);
            }
            sort(multitest.tests[test_index].cells.begin(), multitest.tests[test_index].cells.end());
        }
        return multitest;
    } else if (type == "last_moment") {
        uint32_t lower_n_tests_bound = opt<uint32_t>("lower_n_tests_bound");
        uint32_t upper_n_tests_bound = opt<uint32_t>("upper_n_tests_bound");
        uint32_t lower_m_bound = opt<uint32_t>("lower_m_bound");
        uint32_t upper_m_bound = opt<uint32_t>("upper_m_bound");
        uint32_t lower_x_bound = opt<uint32_t>("lower_coordinate_bound");
        uint32_t upper_x_bound = opt<uint32_t>("upper_coordinate_bound");
        uint32_t lower_total_n_bound = opt<uint32_t>("lower_total_n_bound");
        uint32_t upper_total_n_bound = opt<uint32_t>("upper_total_n_bound");
        int64_t lower_reached_bound = opt<int64_t>("lower_reached_bound");
        int64_t upper_reached_bound = opt<int64_t>("upper_reached_bound");
        uint32_t lower_blue_victories_bound = opt<int32_t>("min_blue_victories");
        uint32_t upper_blue_victories_bound = opt<int32_t>("max_blue_victories");
        upper_n_tests_bound = std::min(upper_n_tests_bound, upper_total_n_bound / 2);
        lower_n_tests_bound = std::max(lower_n_tests_bound, lower_blue_victories_bound);
        uint32_t n_tests = rnd.next(lower_n_tests_bound, upper_n_tests_bound);
        lower_total_n_bound = std::max(lower_total_n_bound, 2 * n_tests);
        uint32_t total_n = rnd.next(lower_total_n_bound, upper_total_n_bound);
        std::vector<uint32_t> test_sizes = decompose_into_sum(total_n - 2 * n_tests, n_tests);
        upper_blue_victories_bound = std::min(upper_blue_victories_bound, n_tests);
        uint32_t blue_victories = rnd.next(lower_blue_victories_bound, upper_blue_victories_bound);
        std::vector<bool> blue_wins(n_tests);
        for (uint32_t test_index = 0; test_index < blue_victories; ++test_index)
            blue_wins[test_index] = true;
        shuffle(blue_wins.begin(), blue_wins.end());
        for (uint32_t& test_size : test_sizes)
            test_size += 2;
        multitest.tests.resize(n_tests);
        for (uint32_t test_index = 0; test_index < n_tests; ++test_index) {
            multitest.tests[test_index].cells.resize(test_sizes[test_index]);
            int64_t reached = rnd.next(lower_reached_bound, upper_reached_bound);
            int32_t ri = 0, bi = 0;
            bool bd = false, rd = false;
            int64_t current_red_minus_blue_sum = 0;
            int64_t total_red_minus_blue_sum = 0;
            int64_t current_red_minus_blue_number = 0;
            int64_t current_coordinate = upper_x_bound;
            /*for (uint32_t cell_index = 0; cell_index < test_sizes[test_index]; ++cell_index) {
                uint32_t remaining_cells = test_sizes[test_index] - cell_index;
                // first things first, let us choose the current coordinate
                while (true) {
                    // if we want to stop, we stop
                    if (rnd.next(current_coordinate - lower_x_bound + 1) < remaining_cells)
                        break;
                    int64_t new_coordinate = current_coordinate - 1; // we can speed up this cycle if here we jump faster, for now it's O(x)
                    // if we cannot make the sum greater than `reached`, we halt
                    if (current_red_minus_blue_number < 0) {
                        // what is the maximum sum we can reach?
                        int64_t needed_cells;
                        if (needed_cells > remaining_cells)
                            break;
                        int64_t maximum_reachable_sum;
                        if (maximum_reachable_sum < reached)
                            break;
                        if (needed_cells < remaining_cells && maximum_reachable_sum == reached)
                            break;
                    }
                    // if we cannot make the sum equal to `reached` in time, we halt
                    current_coordinate = new_coordinate;
                }
            }*/
            while (true) {
                uint32_t red_cells = rnd.next(static_cast<uint32_t>(2), test_sizes[test_index] - 2);
                uint32_t big_red_cells = rnd.next(static_cast<uint32_t>(1), red_cells - 1);
                uint32_t small_red_cells = red_cells - big_red_cells;
                uint32_t blue_cells = test_sizes[test_index] - red_cells;
                uint32_t small_red_cell_l = lower_x_bound, small_red_cell_r = lower_x_bound + small_red_cells - 1;
                uint32_t big_red_cell_r = upper_x_bound, big_red_cell_l = big_red_cell_r - big_red_cells + 1;
                uint64_t total_red_coordinate = 0, total_red_chips = 0;
                for (uint32_t red_cell_index = 0; red_cell_index < red_cells; ++red_cell_index) {
                    if (red_cell_index < small_red_cells) {
                        uint32_t small_red_cell_index = red_cell_index;
                        multitest.tests[test_index].cells[red_cell_index].coordinate = small_red_cell_index + lower_x_bound;
                    } else {
                        uint32_t big_red_cell_index = red_cell_index - small_red_cells;
                        multitest.tests[test_index].cells[red_cell_index].coordinate = upper_x_bound - big_red_cell_index;
                    }
                    multitest.tests[test_index].cells[red_cell_index].color = RED;
                    multitest.tests[test_index].cells[red_cell_index].number_of_chips = rnd.next(lower_m_bound, upper_m_bound);
                    uint64_t number_of_chips = multitest.tests[test_index].cells[red_cell_index].number_of_chips;
                    total_red_chips += number_of_chips;
                    if (red_cell_index == 0)
                        --number_of_chips;
                    total_red_coordinate += number_of_chips * multitest.tests[test_index].cells[red_cell_index].coordinate;
                }
                uint32_t compressed_blue_cells = blue_cells - 1;
                uint64_t total_blue_chips = total_red_chips - 1;
                uint64_t total_blue_chips_temp = total_blue_chips;
                uint64_t total_blue_coordinate = 0;
                std::vector<uint32_t> blue_cell_sizes;
                if (static_cast<uint64_t>(compressed_blue_cells) * upper_m_bound < total_blue_chips || static_cast<uint64_t>(compressed_blue_cells) * lower_m_bound > total_blue_chips)
                    continue;
                for (uint32_t blue_cell_index = 0; blue_cell_index < compressed_blue_cells; ++blue_cell_index) {
                    blue_cell_sizes.push_back(rnd.next(std::max<int64_t>(lower_m_bound, total_blue_chips_temp - static_cast<int64_t>(compressed_blue_cells - 1 - blue_cell_index) * upper_m_bound), std::min<int64_t>(upper_m_bound, total_blue_chips_temp - static_cast<int64_t>(compressed_blue_cells - 1 - blue_cell_index) * lower_m_bound)));
                    total_blue_chips_temp -= blue_cell_sizes.back();
                    total_blue_coordinate += static_cast<uint64_t>(blue_cell_sizes.back()) * blue_cell_index;
                }
                // how to put it?
                // first of all, move it to the right according to the difference between trc and tbc
                int64_t required_total_blue_coordinate = total_red_coordinate - reached;
                if (total_blue_coordinate > required_total_blue_coordinate)
                    continue;
                uint64_t right_translation = (required_total_blue_coordinate - total_blue_coordinate) / total_blue_chips;
                uint64_t blue_cell_l = right_translation, blue_cell_r = blue_cell_l + blue_cells - 1;
                if (blue_cell_l <= small_red_cell_r || blue_cell_r >= big_red_cell_l)
                    continue;
                uint64_t remaining_addend = required_total_blue_coordinate - total_blue_coordinate - total_blue_chips * right_translation;
                // how to add this addend? simply from right to left
                blue_cell_sizes.push_back(0);
                for (int32_t blue_cell_index = compressed_blue_cells - 1; blue_cell_index >= 0; --blue_cell_index) {
                    uint64_t move_to_the_right = std::min<int64_t>(remaining_addend, blue_cell_sizes[blue_cell_index]);
                    blue_cell_sizes[blue_cell_index] -= move_to_the_right;
                    blue_cell_sizes[blue_cell_index + 1] += move_to_the_right;
                    remaining_addend -= move_to_the_right;
                }
                bool has_blue_zero = false;
                for (uint32_t blue_cell_index = 0; blue_cell_index < blue_cells; ++blue_cell_index)
                    if (blue_cell_sizes[blue_cell_index] <= 0) {
                        has_blue_zero = true;
                        break;
                    }
                if (has_blue_zero)
                    continue;
                for (uint32_t blue_cell_index = 0; blue_cell_index < blue_cells; ++blue_cell_index) {
                    multitest.tests[test_index].cells[red_cells + blue_cell_index].coordinate = right_translation + blue_cell_index;
                    multitest.tests[test_index].cells[red_cells + blue_cell_index].number_of_chips = blue_cell_sizes[blue_cell_index];
                    multitest.tests[test_index].cells[red_cells + blue_cell_index].color = BLUE;
                }
                break;
            }
            if (blue_wins[test_index]) {
                for (Cell& cell: multitest.tests[test_index].cells)
                    cell.color = cell.color == BLUE ? RED : BLUE;
            }
            sort(multitest.tests[test_index].cells.begin(), multitest.tests[test_index].cells.end());
        }
        return multitest;
    } else if (type == "frequently_near") {
        uint32_t lower_n_tests_bound = opt<uint32_t>("lower_n_tests_bound");
        uint32_t upper_n_tests_bound = opt<uint32_t>("upper_n_tests_bound");
        uint32_t lower_m_bound = opt<uint32_t>("lower_m_bound");
        uint32_t upper_m_bound = opt<uint32_t>("upper_m_bound");
        uint32_t lower_x_bound = opt<uint32_t>("lower_coordinate_bound");
        uint32_t lower_upper_x_bound = opt<uint32_t>("lower_upper_coordinate_bound");
        uint32_t upper_upper_x_bound = opt<uint32_t>("upper_upper_coordinate_bound");
        uint32_t lower_total_n_bound = opt<uint32_t>("lower_total_n_bound");
        uint32_t upper_total_n_bound = opt<uint32_t>("upper_total_n_bound");
        int64_t lower_reached_bound = opt<int64_t>("lower_reached_bound");
        int64_t upper_reached_bound = opt<int64_t>("upper_reached_bound");
        int64_t lower_frequent_bound = opt<int64_t>("lower_frequent_bound");
        int64_t upper_frequent_bound = opt<int64_t>("upper_frequent_bound");
        uint32_t lower_blue_victories_bound = opt<int32_t>("min_blue_victories");
        uint32_t upper_blue_victories_bound = opt<int32_t>("max_blue_victories");
        upper_n_tests_bound = std::min(upper_n_tests_bound, upper_total_n_bound / 2);
        lower_n_tests_bound = std::max(lower_n_tests_bound, lower_blue_victories_bound);
        uint32_t n_tests = rnd.next(lower_n_tests_bound, upper_n_tests_bound);
        lower_total_n_bound = std::max(lower_total_n_bound, 2 * n_tests);
        uint32_t total_n = rnd.next(lower_total_n_bound, upper_total_n_bound);
        std::vector<uint32_t> test_sizes = decompose_into_sum(total_n - 2 * n_tests, n_tests);
        upper_blue_victories_bound = std::min(upper_blue_victories_bound, n_tests);
        uint32_t blue_victories = rnd.next(lower_blue_victories_bound, upper_blue_victories_bound);
        std::vector<bool> blue_wins(n_tests);
        for (uint32_t test_index = 0; test_index < blue_victories; ++test_index)
            blue_wins[test_index] = true;
        shuffle(blue_wins.begin(), blue_wins.end());
        for (uint32_t& test_size : test_sizes)
            test_size += 2;
        multitest.tests.resize(n_tests);
        for (uint32_t test_index = 0; test_index < n_tests; ++test_index) {
            multitest.tests[test_index].cells.resize(test_sizes[test_index]);
            int64_t reached = rnd.next(lower_reached_bound, upper_reached_bound);
            int64_t frequent = rnd.next(lower_frequent_bound, upper_frequent_bound);
            int32_t current_index = 0;
            uint32_t current_coordinate = rnd.next(lower_upper_x_bound, upper_upper_x_bound);
            auto add = [&](int64_t difference) {
                while (difference) {
                    uint32_t delta, quantity;
                    int64_t coefficient = difference > 0 ? -1 : 1;
                    if (abs(difference) >= upper_m_bound) {
                        delta = abs(difference) / upper_m_bound;
                        quantity = upper_m_bound;
                    } else {
                        delta = 1;
                        quantity = abs(difference);
                    }
                    Color front = difference > 0 ? RED : BLUE;
                    Color back = difference <= 0 ? RED : BLUE;
                    difference += coefficient * delta * quantity;
                    multitest.tests[test_index].cells[current_index] = {
                        .coordinate = current_coordinate,
                        .number_of_chips = quantity,
                        .color = front,
                    };
                    ++current_index;
                    current_coordinate -= delta;
                    multitest.tests[test_index].cells[current_index] = {
                        .coordinate = current_coordinate,
                        .number_of_chips = quantity,
                        .color = back,
                    };
                    ++current_index;
                    --current_coordinate;
                }
            };
            add(frequent);
            while (current_index + 7 <= test_sizes[test_index]) {
                add(upper_m_bound);
                add(-static_cast<int64_t>(upper_m_bound));
            }
            add(reached - frequent);
            uint32_t current_low_x = lower_x_bound;
            while (current_index < test_sizes[test_index]) {
                multitest.tests[test_index].cells[current_index] = {
                    .coordinate = current_low_x,
                    .number_of_chips = upper_m_bound,
                    .color = RED,
                };
                ++current_index;
                ++current_low_x;
            }
            if (blue_wins[test_index]) {
                for (Cell& cell: multitest.tests[test_index].cells)
                    cell.color = cell.color == BLUE ? RED : BLUE;
            }
            sort(multitest.tests[test_index].cells.begin(), multitest.tests[test_index].cells.end());
        }
        return multitest;
    } else {
        ensuref(false, "%s", std::format(R"(Unknown test type: "{}")", type).c_str());
        return {};
    }
}

int main(int argc, char* argv[]) {
    registerGen(argc, argv, 1);
    MultiTest multitest = generate_test();
    std::cout << multitest;
}
