#include <bitset>
#include <span>
#include <cstdint>
#include <iostream>
#include <string>
#include <algorithm>
#include <vector>
#include "VoxelSpace.h"

namespace SomaSolve {
    using SomaSolution = std::vector<uint64_t>;

    struct Solver {
        std::vector<uint64_t>* input;
        std::vector<int>* offsets;
        std::vector<SomaSolution>* solutions;
    };

    auto STD_SOMA = std::vector<uint64_t>{
        23ul,
        30ul,
        15ul,
        1043ul,
        24594ul,
        12306ul,
        11ul, 
    };

    auto backtrack_solve_iter(std::vector<uint64_t> *polycube_input, std::vector<int> *offsets)-> void {
        auto num_inputs = offsets->size() - 1;

        auto solns = std::vector<int>();

        auto iter_stack = std::vector<int>();
        auto curr_soln_stack = std::vector<int>();
        auto soln_spaces_stack = std::vector<uint64_t>();
        soln_spaces_stack.push_back(0ul);

        auto depth = 0;

        while (depth >= 0) {
            if (depth >= iter_stack.size()) {
                iter_stack.push_back(offsets->at(depth));
            }
            auto end = offsets->at(depth + 1);
            auto broke = false;
            for (; iter_stack[depth] < end; iter_stack[depth]++) {
                auto next_space = polycube_input->at(iter_stack[depth]);
                auto soln_space = soln_spaces_stack[depth];
                std::cout << next_space << " " << soln_space << std::endl;
                auto successful_fuse = (soln_space | next_space) == (soln_space ^ next_space);
                if (successful_fuse) {
                    soln_spaces_stack.push_back(soln_space |= next_space);
                    curr_soln_stack.push_back(iter_stack[depth]);
                    depth++;
                    if (curr_soln_stack.size() == num_inputs) {
                        solns.push_back(1);
                        curr_soln_stack.pop_back();
                        soln_spaces_stack.pop_back();
                        depth--;
                    } else {
                        depth++;
                        auto broke = true;
                        break;
                    }
                }
            }
            if (!broke) {
                curr_soln_stack.pop_back();
                soln_spaces_stack.pop_back();
                depth--;
            }
        }
        std::cout << "Done. Found " << solns.size() << " solutions." << std::endl;
    }

    auto backtrack_solve(Solver *solver, uint64_t working_solution = 0ul, int curr_piece = 0) -> void {
        auto input = solver->input;
        auto offsets = solver->offsets;
        auto solutions = solver->solutions;
        auto start = offsets->at(curr_piece);
        auto end = offsets->at(curr_piece + 1);
        auto num_pieces = offsets->size() - 1;
        for (int i = start; i < end; i++) {
            auto successful_fuse = !Voxel::collides(working_solution, input->at(i));
            if (successful_fuse) {
                auto new_working_solution = working_solution | input->at(i);
                solutions->back().at(curr_piece) = input->at(i);
                if (curr_piece == num_pieces - 1) {
                    auto last_soln = solutions->back();
                    solutions->push_back(SomaSolution(last_soln.begin(), last_soln.end()));
                    return;
                } else {
                    backtrack_solve(solver, new_working_solution, curr_piece + 1);
                }
            }
        }
        if (curr_piece == 0) {
            solutions->pop_back();
        } 
    }

    auto get_solution_rotations(SomaSolution *solution, int dims[3]) -> std::vector<SomaSolution> {
        auto result = std::vector<SomaSolution>(Voxel::NUM_ROTS_3D);
        for (int piece_i = 0; piece_i < solution->size(); piece_i++) {
            auto space = Voxel::Space{ 
                .space=solution->at(piece_i),
                .dim_x=dims[0],
                .dim_y=dims[1],
                .dim_z=dims[2],
            };
            auto piece_rotations = Voxel::getAllRotations(&space);
            for (int rot_i = 0; rot_i < piece_rotations.size(); rot_i++) {
                result[rot_i].push_back(piece_rotations[rot_i].space);
            }
        }
        return result;
    }

    auto filter_unique(std::vector<SomaSolution> *solutions, int dims[3]) -> std::vector<SomaSolution> {
        if (solutions->size() == 0) {
            return std::vector<SomaSolution>();
        }
        auto unique_solns = std::vector<SomaSolution>{};
        for (auto &solution : *solutions) {
            auto found_match = false;
            for (auto &rotation : get_solution_rotations(&solution, dims)) { 
                for (auto &unique_soln : unique_solns) {
                    auto is_match = true;
                    for (int piece_i = 0; piece_i < unique_soln.size(); piece_i++) {
                        if (rotation[piece_i] != unique_soln[piece_i]) {
                            is_match = false;
                            break;
                        }
                    }
                    if (is_match) {
                        found_match = true;
                        break;
                    }
                }
                if (found_match) {
                    break;
                }
            }
            if (!found_match) {
                unique_solns.push_back(SomaSolution(solution));
            }
        }
        return unique_solns;
    }

    auto solve(std::vector<uint64_t> *reprs_in, int dims[3]) -> std::vector<SomaSolution> {
        auto reprs = *reprs_in;
        auto offsets = std::vector<int>();
        auto polycubes = std::vector<uint64_t>();
        polycubes.reserve(reprs.size() * 10);

        auto model_space = Voxel::Space{ 
            .space={},
            .dim_x=dims[0],
            .dim_y=dims[1],
            .dim_z=dims[2],
        };

        offsets.push_back(0);
        auto space = model_space;
        space.space = reprs[0];
        Voxel::cullEmptySpace(&space);
        auto positions = Voxel::getAllPositionsInPrism(&space, dims);
        polycubes.insert(polycubes.end(), positions.begin(), positions.end());

        for (int i = 1; i < reprs.size(); i++) {
            offsets.push_back(polycubes.size());
            auto space = model_space;
            space.space = reprs[i];
            Voxel::cullEmptySpace(&space);
            auto perms = Voxel::getAllPermutationsInPrism(&space, dims);
            polycubes.insert(polycubes.end(), perms.begin(), perms.end());
        }

        offsets.push_back(polycubes.size());

        auto solutions = std::vector<SomaSolution>{std::vector<uint64_t>(reprs.size())};
        auto solver = Solver{
            .input=&polycubes,
            .offsets=&offsets,
            .solutions=&solutions,
        };

        backtrack_solve(&solver);

        return filter_unique(solver.solutions, dims);
    }
}