#pragma once

#include <tc/core.hpp>
#include <cmath>
#include <optional>
#include <numeric>
#include <iostream>

#include <geometry.hpp>
#include <utility>

#include "combo_iterator.hpp"

/**
 * Produce a list of all generators for the group context. The range [0..group.ngens).
 */
std::vector<int> generators(const tc::Group &context) {
    std::vector<int> g_gens(context.ngens);
    std::iota(g_gens.begin(), g_gens.end(), 0);
    return g_gens;
}

namespace {
    /**
     * Determine which of g_gens are the correct names for sg_gens within the current context
     */
    std::vector<int> recontext_gens(
        const tc::Group &context,
        std::vector<int> g_gens,
        std::vector<int> sg_gens) {

        std::sort(g_gens.begin(), g_gens.end());

        int inv_gen_map[context.ngens];
        for (size_t i = 0; i < g_gens.size(); i++) {
            inv_gen_map[g_gens[i]] = i;
        }

        std::vector<int> s_sg_gens;
        s_sg_gens.reserve(sg_gens.size());
        for (const auto gen : sg_gens) {
            s_sg_gens.push_back(inv_gen_map[gen]);
        }
        std::sort(s_sg_gens.begin(), s_sg_gens.end());

        return s_sg_gens;
    }

    /**
     * Solve the cosets generated by sg_gens within the subgroup generated by g_gens of the group context
     */
    tc::Cosets solve(
        const tc::Group &context,
        const std::vector<int> &g_gens,
        const std::vector<int> &sg_gens
    ) {
        const auto proper_sg_gens = recontext_gens(context, g_gens, sg_gens);
        return context.subgroup(g_gens).solve(proper_sg_gens);
    }

    /**
     * Apply some context transformation to all primitives of this mesh.
     */
    template<unsigned N>
    void apply(const tc::Cosets &table, int gen, Prims<N> &mat) {
        auto data = mat.data();
        for (int i = 0; i < mat.size(); ++i) {
            data[i] = table.get(data[i], gen);
        }
    }
}

template<unsigned N>
class Mesh {
public:
    const tc::Group *g;  // todo this needs to be handled more consistently
    std::vector<int> ctx;
    Prims<N> prims;

    Mesh(const tc::Group &g_, std::vector<int> ctx_, size_t cols);

    static Mesh<N> fill(const tc::Group &g, std::vector<int> ctx);

    template<class SC>
    static Mesh<N> hull(const tc::Group &g, std::vector<int> ctx, const SC &sub_ctxs);

    Mesh<N> recontext(std::vector<int> ctx_);

    std::vector<Mesh<N>> tile(const std::vector<int> &ctx_);

    Mesh<N + 1> fan(unsigned root);

    [[nodiscard]] size_t size() const { return prims.size(); }

    [[nodiscard]] size_t rows() const { return prims.rows(); }

    [[nodiscard]] size_t cols() const { return prims.cols(); }

    [[nodiscard]] unsigned *data() { return prims.data(); }

    [[nodiscard]] const unsigned *data() const { return prims.data(); }
};

template<class M>
M merge(const std::vector<M> &meshes) {
    if (meshes.empty()) throw std::logic_error("cannot merge an empty list of meshes");

    auto g = meshes[0].g;
    auto ctx = meshes[0].ctx;

    size_t cols = 0;
    for (const auto &mesh : meshes) {
        cols += mesh.prims.cols();
    }

    M res(*g, ctx, cols);

    size_t offset = 0;
    for (const auto &mesh : meshes) {
        res.prims.middleCols(offset, mesh.prims.cols()) = mesh.prims;
        offset += mesh.prims.cols();
    }

    return res;
}

template<unsigned N>
Mesh<N> Mesh<N>::recontext(std::vector<int> ctx_) {
    Mesh<N> res = *this;
    res.ctx = ctx_;

    const auto proper_sg_gens = recontext_gens(*g, res.ctx, ctx);
    const auto table = solve(*g, res.ctx, {});
    const auto path = solve(*g, ctx, {}).path;

    auto map = path.template walk<int, int>(0, proper_sg_gens, [table](int coset, int gen) {
        return table.get(coset, gen);
    });

    auto data = res.prims.data();
    for (int i = 0; i < res.prims.size(); ++i) {
        data[i] = map[data[i]];
    }

    return res;
}

template<unsigned N>
std::vector<Mesh<N>> Mesh<N>::tile(const std::vector<int> &ctx_) {
    auto base = recontext(ctx_);

    auto table = solve(*g, base.ctx, {});
    auto path = solve(*g, base.ctx, ctx).path;

    std::vector<Mesh<N>> res = path.template walk<Mesh<N>, int>(
        base, generators(*g),
        [&](Mesh<N> mesh, int gen) {
            apply<N>(table, gen, mesh.prims);
            return mesh;
        }
    );

    return res;
}

template<unsigned N>
Mesh<N + 1> Mesh<N>::fan(unsigned root) {
    Mesh<N + 1> res(*g, ctx, prims.cols());

    res.prims.topRows(1) = Prims<1>::Constant(1, prims.cols(), root);
    res.prims.bottomRows(N) = prims;

    return res;
}

template<unsigned N>
Mesh<N>::Mesh(const tc::Group &g_, std::vector<int> ctx_, size_t cols)
    : g(&g_), ctx(std::move(ctx_)) {
    prims.setZero(N, cols);
}

template<unsigned N>
Mesh<N> Mesh<N>::fill(const tc::Group &g, std::vector<int> ctx) {
    if (ctx.size() + 1 != N)
        throw std::logic_error("ctx size must be one less than N");

    const auto &combos = Combos(ctx, (int)ctx.size() - 1);

    std::vector<Mesh<N>> meshes;

    for (const auto &sub_ctx : combos) {
        auto base = Mesh<N - 1>::fill(g, sub_ctx);
        auto parts = base.tile(ctx);
        parts.erase(parts.begin(), parts.begin() + 1);

        if (parts.empty()) continue;
        auto raised = merge(parts);

        auto fanned = raised.fan(0);
        meshes.push_back(fanned);
    }

    return merge(meshes);
}

template<>
Mesh<1> Mesh<1>::fill(const tc::Group &g, std::vector<int> ctx) {
    if (not ctx.empty())
        throw std::logic_error("ctx must be empty for a trivial Mesh.");

    return Mesh<1>(g, ctx, 1);
}

template<unsigned N>
template<class SC>
Mesh<N> Mesh<N>::hull(const tc::Group &g, const std::vector<int> ctx, const SC &sub_ctxs) {
    std::vector<Mesh<N>> parts;

    for (const auto &sub_ctx : sub_ctxs) {
        auto face = Mesh<N>::fill(g, sub_ctx).tile(ctx);
        parts.insert(parts.end(), face.begin(), face.end());
    }

    return merge(parts);
}