#include <cstring>
#include <iostream>
#include <string>
#include <vector>
#include "VoxelSpace.h"
#include "lib/djstdlib/core.h"

void get_dims_input(int dims[3]) {
    std::cout << "Enter dimensions separated by newlines. (x*y*z must not exceed 64)\n";
    bool success = false;
    while (!success) {
        std::cout << "x: ";
        std::cin >> dims[0];
        std::cout << "y: ";
        std::cin >> dims[1];
        std::cout << "z: ";
        std::cin >> dims[2];

        int size = dims[0]*dims[1]*dims[2];
        if (size <= 64) {
            success = true;
        } else {
            std::cout << "That resulted in " << size << " units. Try again.\n"; 
        }
        std::cin.ignore();
    }
}

std::vector<uint64> get_reprs_input(int units_required) {
    std::cout << "Enter bit-representations (big endian, max 64 bits, total 1s must add up to " << units_required << "). press ENTER twice to finish input.\n";
    std::vector<uint64> reprs = std::vector<uint64>();
    int total_units = 0;
    while (true) {
        std::string input = std::string();
        std::getline(std::cin, input);
        if (input.size() == 0) {
            if (total_units == units_required) {
                break;
            } else {
                std::cout << "Bad number of units. You entered: " << total_units << ", but exactly " << units_required << " were required.\n";
                total_units = 0;
                continue;
            }
        }
        uint64 bit_repr = 0;
        int i = 0;
        bool good_repr = true;
        for (auto it = input.rbegin(); it < input.rend(); it++, i++) {
            if (*it == '1') {
                bit_repr |= 1ull << i;
                total_units++;
            } else if (*it != '0' || i >= 64) {
                std::cout << "Input invalid. Enter a binary string only with max 64 bits." << '\n';
                good_repr = false;
                break;
            }
        }
        if (good_repr) {
            reprs.push_back(bit_repr);
        }
    }
    return reprs;
}

typedef std::vector<uint64> SomaSolution;

struct Solver {
    list<uint64>* input;
    list<size_t>* offsets;
    std::vector<SomaSolution>* solutions;
};

uint64 STD_SOMA[] = { 23ul, 30ul, 15ul, 1043ul, 24594ul, 12306ul, 11ul };

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

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

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

    int depth = 0;

    while (depth >= 0) {
        if (depth >= iter_stack.size()) {
            iter_stack.push_back(offsets->at(depth));
        }
        int end = offsets->at(depth + 1);
        bool broke = false;
        for (; iter_stack[depth] < end; iter_stack[depth]++) {
            uint64 next_space = polycube_input->at(iter_stack[depth]);
            uint64 soln_space = soln_spaces_stack[depth];
            std::cout << next_space << " " << soln_space << std::endl;
            bool 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++;
                    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;
}

void backtrack_solve(Solver *solver, uint64 working_solution = 0, size_t curr_piece = 0) {
    list<uint64> *input = solver->input;
    list<size_t> *offsets = solver->offsets;
    std::vector<SomaSolution> *solutions = solver->solutions;
    size_t start = offsets->data[curr_piece];
    size_t end = offsets->data[curr_piece + 1];
    size_t num_pieces = offsets->length - 1;
    for (size_t i = start; i < end; i++) {
        bool successful_fuse = !collides(working_solution, input->data[i]);
        if (successful_fuse) {
            uint64 new_working_solution = working_solution | input->data[i];
            solutions->back().at(curr_piece) = input->data[i];
            if (curr_piece == num_pieces - 1) {
                std::vector<uint64> 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();
    } 
}

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

std::vector<SomaSolution> filter_unique(Arena *arena, std::vector<SomaSolution> *solutions, int dims[3]) {
    if (solutions->size() == 0) {
        return std::vector<SomaSolution>();
    }
    std::vector<SomaSolution> unique_solns = std::vector<SomaSolution>{};
    for (std::vector<uint64> &solution : *solutions) {
        bool found_match = false;
        Scratch temp = scratchStart(&arena, 1);
        std::vector<SomaSolution> rots = get_solution_rotations(temp.arena, &solution, dims);
        for (SomaSolution &rotation : rots) { 
            for (std::vector<uint64> &unique_soln : unique_solns) {
                bool 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;
            }
        }
        scratchEnd(temp);
        if (!found_match) {
            unique_solns.push_back(SomaSolution(solution));
        }
    }
    return unique_solns;
}

std::vector<SomaSolution> solve(uint64 *reprs_in, uint32 reprs_in_count, int dims[3]) {
    Arena *arena = arenaAlloc(Megabytes(64));
    Arena *permsArena = arenaAlloc(Megabytes(64));

    list<size_t> offsets = PushList(arena, size_t, reprs_in_count + 1);

    list<uint64> polycubes = PushList(arena, uint64, 0);

    Space empty_voxel_space = { 
        {},
        dims[0],
        dims[1],
        dims[2],
    };

    appendList(&offsets, (size_t)0);

    list<uint64> positions = {};
    {
        Space space = empty_voxel_space;
        space.space = reprs_in[0];
        cullEmptySpace(&space);
        positions = getAllPositionsInPrism(permsArena, &space, dims);
        uint64 *insertion = PushArray(arena, uint64, positions.length);
        polycubes.length += positions.length;
        polycubes.head += positions.head;
        memcpy(insertion, positions.data, positions.length * sizeof(uint64));
    };

    for (size_t i = 1; i < reprs_in_count; i++) {
        appendList(&offsets, polycubes.length);
        Space space = empty_voxel_space;
        space.space = reprs_in[i];
        cullEmptySpace(&space);
        list<uint64> perms = getAllPermutationsInPrism(permsArena, &space, dims);
        uint64 *insertion = PushArray(arena, uint64, perms.length);
        polycubes.length += perms.length;
        polycubes.head += perms.head;
        memcpy(insertion, perms.data, perms.length * sizeof(uint64));
    }

    appendList(&offsets, polycubes.length);

    std::vector<SomaSolution> solutions = {std::vector<uint64>(reprs_in_count)};
    Solver solver = {
        &polycubes,
        &offsets,
        &solutions,
    };

    backtrack_solve(&solver);

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


void interactive_cmd_line_solve_soma() {
    int dims[3] = { 3, 3, 3 };
    //get_dims_input(dims);
    //std::cout << '\n';
    //std::vector<uint64> reprs = get_reprs_input(dims[0]*dims[1]*dims[2]);
    std::cout << "Great. Calculating solutions...\n";
    std::vector<SomaSolution> solutions = solve(STD_SOMA, ArrayCount(STD_SOMA), dims);
    std::cout << solutions.size() << " solutions found." << std::endl;
}