#pragma once

#include <vector>
#include <cassert>
#include <limits>
#include <algorithm>
#include <numeric>
#include <functional>

#include <oineus/filtration.h>

namespace std {
template<typename T>
std::ostream& operator<<(std::ostream& os, const std::vector<T>& v)
{
    os << "[";

    for(size_t i = 0 ; i < v.size() ; ++i) {
        os << v[i];
        if (i + 1 < v.size())
            os << ", ";
    }

    os << "]";
    return os;
}

template<typename Int, size_t D>
std::ostream& operator<<(std::ostream& os, const std::array<Int, D>& p)
{
    os << "(";

    for(size_t i = 0 ; i < D ; ++i) {
        os << p[i];
        if (i + 1 < D)
            os << ", ";
    }

    os << ")";
    return os;
}
}

namespace oineus {

// this function is auto-generated by scripts/freudenthal_cube.py
template<typename Int, size_t D>
constexpr inline std::vector<std::vector<std::vector<Int>>> fr_displacements(size_t dim)
{
    size_t cube_dim = D;

    switch(cube_dim) {
    case 1 :
        switch(dim) {
        case 0 :
            return {
                    {{0}}};
            break;
        case 1 :
            return {
                    {{0}, {1}}};
            break;
        }
        break;
    case 2 :
        switch(dim) {
        case 0 :
            return {
                    {{0, 0}}};
            break;
        case 1 :
            return {
                    {{0, 0}, {0, 1}},
                    {{0, 0}, {1, 1}},
                    {{0, 0}, {1, 0}}};
            break;
        case 2 :
            return {
                    {{0, 0}, {1, 0}, {1, 1}},
                    {{0, 0}, {0, 1}, {1, 1}}};
            break;
        }
        break;
    case 3 :
        switch(dim) {
        case 0 :
            return {
                    {{0, 0, 0}}};
            break;
        case 1 :
            return {
                    {{0, 0, 0}, {0, 1, 0}},
                    {{0, 0, 0}, {1, 1, 0}},
                    {{0, 0, 0}, {0, 1, 1}},
                    {{0, 0, 0}, {1, 1, 1}},
                    {{0, 0, 0}, {0, 0, 1}},
                    {{0, 0, 0}, {1, 0, 0}},
                    {{0, 0, 0}, {1, 0, 1}}};
            break;
        case 2 :
            return {
                    {{0, 0, 0}, {0, 0, 1}, {0, 1, 1}},
                    {{0, 0, 0}, {0, 0, 1}, {1, 0, 1}},
                    {{0, 0, 0}, {1, 0, 0}, {1, 0, 1}},
                    {{0, 0, 0}, {1, 0, 0}, {1, 1, 1}},
                    {{0, 0, 0}, {0, 1, 0}, {1, 1, 0}},
                    {{0, 0, 0}, {1, 0, 1}, {1, 1, 1}},
                    {{0, 0, 0}, {0, 1, 1}, {1, 1, 1}},
                    {{0, 0, 0}, {0, 1, 0}, {0, 1, 1}},
                    {{0, 0, 0}, {0, 0, 1}, {1, 1, 1}},
                    {{0, 0, 0}, {1, 0, 0}, {1, 1, 0}},
                    {{0, 0, 0}, {0, 1, 0}, {1, 1, 1}},
                    {{0, 0, 0}, {1, 1, 0}, {1, 1, 1}}};
            break;
        case 3 :
            return {
                    {{0, 0, 0}, {0, 0, 1}, {0, 1, 1}, {1, 1, 1}},
                    {{0, 0, 0}, {1, 0, 0}, {1, 0, 1}, {1, 1, 1}},
                    {{0, 0, 0}, {0, 1, 0}, {1, 1, 0}, {1, 1, 1}},
                    {{0, 0, 0}, {0, 0, 1}, {1, 0, 1}, {1, 1, 1}},
                    {{0, 0, 0}, {1, 0, 0}, {1, 1, 0}, {1, 1, 1}},
                    {{0, 0, 0}, {0, 1, 0}, {0, 1, 1}, {1, 1, 1}}};
            break;
        }
        break;
    }

    throw std::runtime_error("Unknown dimensions");
}

template<typename Int_, typename Real_, size_t D>
class Grid {
public:
    using Int = Int_;
    using Real = Real_;
    using GridSimplex = CellWithValue<Simplex<Int>, Real>;
    using GridPoint = std::array<Int, D>;
    using GridPointVec = std::vector<GridPoint>;
    using GridPointVecVec = std::vector<GridPointVec>;
    using SimplexVec = std::vector<GridSimplex>;
    using CriticalVertices = std::vector<Int>;
    using IdxVector = typename Simplex<Int>::IdxVector;

    using GridFiltration = Filtration<Simplex<Int>, Real>;

    static constexpr size_t dim {D};

    struct ValueVertex {
        Real value;
        Int vertex;
    };

    Grid() = default;

    Grid(const GridPoint& _dims, bool _wrap, Real* _data)
            :dims_(_dims), c_order_(true), wrap_(_wrap), data_(_data)
    {
        if (dims_.size() != dim)
            throw std::runtime_error("Dims size must be equal D");

        if (wrap_ and std::any_of(dims_.begin(), dims_.end(), [](Int x) { return x < 3; }))
            throw std::runtime_error("Must have at least 3 cells in each dimension for wrap=true");

        if (not std::all_of(dims_.begin(), dims_.end(), [](Int x) { return x > 0; }))
            throw std::runtime_error("Dims must be positive");

        if (c_order_) {
            strides_[D - 1] = 1;

            for(int d = static_cast<int>(D) - 2 ; d >= 0 ; --d)
                strides_[d] = dims_[d + 1] * strides_[d + 1];
        } else
            throw std::runtime_error("Fortran order not implemented");
    }

    Int size() const
    {
        return std::accumulate(dims_.cbegin(), dims_.cend(), Int(1), std::multiplies<Int>());
    }

    Int point_to_id(const GridPoint& v) const
    {
        assert(c_order_);
        return std::inner_product(v.begin(), v.end(), strides_.begin(), 0);
    }

    GridPoint wrap_point(const GridPoint& v) const
    {
        assert(wrap_);
        GridPoint result = v;
        for(size_t i = 0 ; i < dim ; ++i)
            result[i] %= dims_[i];
        return result;
    }

    bool point_in_bounds(const GridPoint& v) const
    {
        for(size_t i = 0 ; i < dim ; ++i)
            if (v[i] >= dims_[i] or v[i] < 0)
                return false;
        return true;
    }

    GridPoint id_to_point(Int i) const
    {
        GridPoint result;

        assert(0 <= i and i < size());

        if (c_order_)
            for(size_t d = 0 ; d < dim ; ++d) {
                result[d] = i / strides_[d];
                i = i % strides_[d];
            }
        else
            throw std::runtime_error("Fortran order not supported");
        return result;
    }

    template<class Cont>
    static GridPoint add_points(const GridPoint& x, const Cont& y)
    {
        assert(x.size() == y.size());

        GridPoint z = x;
        for(size_t i = 0 ; i < dim ; ++i)
            z[i] += y[i];

        return z;
    }

    GridPointVecVec get_fr_displacements(size_t d) const
    {
        GridPointVecVec result;

        for(auto&& disps: fr_displacements<Int, D>(d)) {
            result.emplace_back();
            for(auto&& disp: disps) {
                GridPoint p;
                assert(disp.size() == p.size());
                for(size_t k = 0 ; k < disp.size() ; ++k)
                    p[k] = disp[k];
                result.back().push_back(p);
            }
        }

        return result;
    }

    std::pair<GridFiltration, CriticalVertices> freudenthal_filtration_and_critical_vertices(size_t top_d, bool negate, int n_threads = 1) const
    {
        if (top_d > dim)
            throw std::runtime_error("bad dimension, top_d = " + std::to_string(top_d) + ", dim = " + std::to_string(dim));

        SimplexVec simplices;
        CriticalVertices vertices;

        // calculate total number of cells to allocate memory once
        size_t total_size = 0;
        for(dim_type d = 0 ; d <= top_d ; ++d) {
            total_size += get_fr_displacements(d).size() * size();
        }

        simplices.reserve(total_size);
        vertices.reserve(total_size);

        for(dim_type d = 0 ; d <= top_d ; ++d) {
            add_freudenthal_simplices(d, negate, simplices, vertices, true);
        }

        auto fil = GridFiltration(simplices, negate, n_threads);

        CriticalVertices sorted_vertices;

        for(const auto& cell: fil.cells()) {
            sorted_vertices.push_back(vertices[cell.get_id()]);
        }

        return {fil, sorted_vertices};
    }

    GridFiltration freudenthal_filtration(size_t top_d, bool negate, int n_threads = 1) const
    {
        if (top_d > dim)
            throw std::runtime_error("bad dimension, top_d = " + std::to_string(top_d) + ", dim = " + std::to_string(dim));

        SimplexVec simplices;
        CriticalVertices dummy_vertices;

        // calculate total number of cells to allocate memory once
        size_t total_size = 0;
        for(dim_type d = 0 ; d <= top_d ; ++d) {
            total_size += get_fr_displacements(d).size() * size();
        }

        simplices.reserve(total_size);

        for(dim_type d = 0 ; d <= top_d ; ++d) {
            add_freudenthal_simplices(d, negate, simplices, dummy_vertices, false);
        }

        auto fil = GridFiltration(simplices, negate, n_threads);
        return fil;
    }

    Real value_at_vertex(Int vertex) const
    {
        assert(vertex >= 0 and vertex < size());
        return *(data_ + vertex);
    }

    ValueVertex simplex_value_and_vertex(const IdxVector& vertices, bool negate) const
    {
        ValueVertex result {negate ? std::numeric_limits<Real>::max() : std::numeric_limits<Real>::lowest(), Int(-1)};

        // break ties by vertex id, ignore negate for ties
        auto cmp = [negate](ValueVertex x, ValueVertex y) {
          return negate ? (x.value > y.value or (x.value == y.value and x.vertex < y.vertex)) : (x.value < y.value or (x.value == y.value and x.vertex < y.vertex));
        };

        for(Int v: vertices) {
            ValueVertex cand {value_at_vertex(v), v};
            result = std::max(result, cand, cmp);
        }

        return result;
    }

    template<typename I, typename R, size_t DD>
    friend std::ostream& operator<<(std::ostream& out, const Grid<I, R, DD>& g);

private:
    GridPoint dims_;
    GridPoint strides_;
    bool c_order_ {true};
    bool wrap_ {false};
    Real* data_ {nullptr};

    void add_freudenthal_simplices_from_vertex(const GridPoint& v,
            size_t d,
            bool negate,
            const GridPointVecVec& disps,
            SimplexVec& simplices,
            CriticalVertices& critical_vertices,
            bool return_critical_vertices) const
    {
        IdxVector v_ids(d + 1, 0);

        for(auto& deltas: disps) {
            assert(deltas.size() == d + 1);

            bool is_valid_simplex = true;

            for(size_t i = 0 ; i < d + 1 ; ++i) {
                GridPoint u = add_points(v, deltas[i]);
                if (wrap_)
                    v_ids[i] = point_to_id(wrap_point(u));
                else if (not point_in_bounds(u)) {
                    is_valid_simplex = false;
                    break;
                } else
                    v_ids[i] = point_to_id(u);
            }

            if (is_valid_simplex) {
                ValueVertex vv = simplex_value_and_vertex(v_ids, negate);
                simplices.emplace_back(v_ids, vv.value);
                if (return_critical_vertices) {
                    critical_vertices.emplace_back(vv.vertex);
                }
            }
        }
    }

    void add_freudenthal_simplices(dim_type d,
                                   bool negate,
                                   SimplexVec& simplices,
                                   CriticalVertices& critical_vertices,
                                   bool return_critical_vertices) const
    {
        auto disps = get_fr_displacements(d);

        for(Int i = 0 ; i < size() ; ++i) {
            GridPoint v = id_to_point(i);
            add_freudenthal_simplices_from_vertex(v, d, negate, disps, simplices, critical_vertices, return_critical_vertices);

#ifdef OINEUS_CHECK_FOR_PYTHON_INTERRUPT
            if (i % 100 == 0) {
                OINEUS_CHECK_FOR_PYTHON_INTERRUPT;
            }
#endif

        }
    }
}; // class Grid

template<typename Int, typename Real, size_t D>
std::ostream& operator<<(std::ostream& out, const Grid<Int, Real, D>& g)
{
    out << "Grid(dims=" << g.dims_ << ", strides = " << g.strides_ << ", c_order = " << g.c_order_ << ", wrap = " << g.wrap_ << ", data = " << g.data_ << ")";
    return out;
}

} // namespace oineus
