#pragma once

#include <vector>
#include <array>
#include <gmsh.h>
#include <type_traits>

#ifdef USE_MPI
#include <mpi/mpi.h>
#endif /* USE_MPI */

#include "entity.h"
#include "physical_element.h"

namespace g     = gmsh;
namespace go    = gmsh::option;
namespace gm    = gmsh::model;
namespace gmm   = gmsh::model::mesh;
namespace gmo   = gmsh::model::occ;
namespace gmg   = gmsh::model::geo;
namespace gv    = gmsh::view;

using gvp_t = gmsh::vectorpair;

/* Macros used to annotate GMSH integer inputs */
#define DIMENSION(x) x

std::string quadrature_name(int);
std::string basis_func_name(int);
std::string basis_grad_name(int);

using face_key = std::array<size_t, 3>;

enum class face_type : int
{
    NONE = 0,
    INTERFACE,
    BOUNDARY,
    INTERPROCESS_BOUNDARY
};

struct boundary_descriptor
{
    face_type   type;
    int         gmsh_entity;

    //auto operator <=> (const boundary_descriptor&) const = default;

    bool operator<(const boundary_descriptor& other) const
    {
        return (type < other.type) or
               ( (type == other.type) and (gmsh_entity < other.gmsh_entity) );
    }

    bool operator==(const boundary_descriptor& other) const
    {
        return (type == other.type) and (gmsh_entity == other.gmsh_entity);
    }

    boundary_descriptor()
        : type(face_type::NONE), gmsh_entity(0)
    {}
};

#if ((!defined(NDEBUG)) && defined(USE_MPI))
    #define ASSERT_MPI_RANK_0 {                 \
        int __rank;                             \
        MPI_Comm_rank(MPI_COMM_WORLD, &__rank); \
        assert(__rank == 0);                    \
    }
#else
    #define ASSERT_MPI_RANK_0 (void) 0
#endif

#define IMPLIES(a, b) ((not(a)) or (b))
#define IFF(a,b) (IMPLIES(a,b) and IMPLIES(b,a))

template<typename T>
void
priv_MPI_Bcast(std::vector<T>& vec, int root, MPI_Comm comm)
{
    static_assert(std::is_trivially_copyable<T>::value, "Type must be trivially copyable");

    int rank;
    MPI_Comm_rank(comm, &rank);

    if (rank == root)
    {
        size_t vsize = vec.size();
        MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        MPI_Bcast(vec.data(), vec.size()*sizeof(T), MPI_PACKED, root, comm);
    }
    else
    {
        size_t vsize;
        MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        vec.resize(vsize);
        MPI_Bcast(vec.data(), vec.size()*sizeof(T), MPI_PACKED, root, comm);
    }
}

template<typename T1, typename T2>
void
priv_MPI_Bcast(std::map<T1, std::vector<T2>>& map, int root, MPI_Comm comm)
{
    static_assert(std::is_trivially_copyable<T1>::value, "Type must be trivially copyable");
    static_assert(std::is_trivially_copyable<T2>::value, "Type must be trivially copyable");

    int rank;
    MPI_Comm_rank(comm, &rank);

    if (rank == root)
    {
        size_t msize = map.size();
        MPI_Bcast(&msize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        for (auto& [l, rv] : map)
        {
            auto ll = l;
            MPI_Bcast(&ll, sizeof(T1), MPI_PACKED, root, comm);
            size_t vsize = rv.size();
            MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
            MPI_Bcast(rv.data(), vsize*sizeof(T2), MPI_PACKED, root, comm);
        }
    }
    else
    {
        size_t msize;
        MPI_Bcast(&msize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        for (size_t i = 0; i < msize; i++)
        {
            T1 l;
            MPI_Bcast(&l, sizeof(T1), MPI_PACKED, root, comm);
            size_t vsize;
            MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
            std::vector<T2> rv;
            rv.resize(vsize);
            MPI_Bcast(rv.data(), vsize*sizeof(T2), MPI_PACKED, root, comm);
            map[l] = std::move(rv);
        }
    }
}

class model
{
    int                                         geometric_order;
    int                                         approximation_order;
    int                                         num_partitions;

    std::map<int, std::vector<element_key>>     boundary_map;
    std::map<int, std::vector<int>>             comm_map;
    std::map<int, std::vector<int>>             partition_map;

    std::vector<entity>                         entities;
    std::vector<boundary_descriptor>            bnd_descriptors;
    element_connectivity<element_key>           conn;

    using entofs_pair = std::pair<size_t, size_t>;
    std::map<size_t, entofs_pair>               etag_to_entity_offset;

    void    map_boundaries(void);
    void    import_gmsh_entities(void);
    void    update_connectivity(const entity&, size_t);

    void    populate_from_gmsh_rank0(void);

#ifdef USE_MPI
    void    populate_from_gmsh_rankN(void);
#endif

    void    make_boundary_to_faces_map(void);

public:
    model();
    model(int, int);
    model(const char *);
    model(const char *, int, int);
    ~model();

    void build(void);
    void build(double);

    void generate_mesh(void);
    void generate_mesh(double);
    void partition_mesh(int);
    void populate_from_gmsh(void);

    const element_connectivity<element_key>& connectivity() const
    {
        return conn;
    }

    const std::vector<boundary_descriptor>& boundary_descriptors() const
    {
        return bnd_descriptors;
    }

    std::vector<element_key> get_bnd(size_t which)
    {
        return boundary_map.at(which);
    }

    size_t  num_dofs() const;
    size_t  num_fluxes() const;
    size_t  num_cells() const;
    size_t  num_faces() const;
    size_t  num_entities() const;

    std::vector<entity>::const_iterator begin() const;
    std::vector<entity>::const_iterator end() const;
    std::vector<entity>::iterator begin();
    std::vector<entity>::iterator end();

    entity&         entity_at(size_t);
    const entity&   entity_at(size_t) const;

    entofs_pair     lookup_tag(size_t tag) const;
};