#ifndef MESH_DATA_HPP
#define MESH_DATA_HPP

#include <Kokkos_Core.hpp>
#include <TinyVector.hpp>

#include <SubItemValuePerItem.hpp>

#include <map>

template <typename M>
class MeshData
{
 public:
  typedef M MeshType;

  static constexpr size_t dimension = MeshType::dimension;
  static_assert(dimension>0, "dimension must be strictly positive");

  typedef TinyVector<dimension> Rd;

  static constexpr double inv_dimension = 1./dimension;

 private:
  const MeshType& m_mesh;
  NodeValuePerCell<Rd> m_Cjr;
  NodeValuePerCell<double> m_ljr;
  NodeValuePerCell<Rd> m_njr;
  Kokkos::View<Rd*>  m_xj;
  Kokkos::View<double*>   m_Vj;

  KOKKOS_INLINE_FUNCTION
  void _updateCenter()
  { // Computes vertices isobarycenter
    if constexpr (dimension == 1) {
      const Kokkos::View<const Rd*> xr = m_mesh.xr();

      Kokkos::parallel_for(m_mesh.numberOfCells(), KOKKOS_LAMBDA(const int& j){
          const auto& cell_nodes = m_mesh.connectivity().m_cell_to_node_matrix.rowConst(j);
          m_xj[j] = 0.5*(xr[cell_nodes(0)]+xr[cell_nodes(1)]);
        });
    } else {
      const Kokkos::View<const Rd*> xr = m_mesh.xr();
      const Kokkos::View<const double*>& inv_cell_nb_nodes
          = m_mesh.connectivity().invCellNbNodes();
      Kokkos::parallel_for(m_mesh.numberOfCells(), KOKKOS_LAMBDA(const int& j){
          Rd X = zero;
          const auto& cell_nodes = m_mesh.connectivity().m_cell_to_node_matrix.rowConst(j);
          for (size_t R=0; R<cell_nodes.length; ++R) {
            X += xr[cell_nodes(R)];
          }
          m_xj[j] = inv_cell_nb_nodes[j]*X;
        });
    }
  }

  KOKKOS_INLINE_FUNCTION
  void _updateVolume()
  {
    const Kokkos::View<const Rd*> xr = m_mesh.xr();

    Kokkos::parallel_for(m_mesh.numberOfCells(), KOKKOS_LAMBDA(const int& j){
        double sum_cjr_xr = 0;
        const auto& cell_nodes = m_mesh.connectivity().m_cell_to_node_matrix.rowConst(j);

        for (size_t R=0; R<cell_nodes.length; ++R) {
          sum_cjr_xr += (xr[cell_nodes(R)], m_Cjr(j,R));
        }
        m_Vj[j] = inv_dimension * sum_cjr_xr;
      });
  }

  KOKKOS_INLINE_FUNCTION
  void _updateCjr() {
    if constexpr (dimension == 1) {
      // Cjr/njr/ljr are constant overtime
    }
    else if constexpr (dimension == 2) {
      const Kokkos::View<const Rd*> xr = m_mesh.xr();

      Kokkos::parallel_for(m_mesh.numberOfCells(), KOKKOS_LAMBDA(const int& j){
          const auto& cell_nodes = m_mesh.connectivity().m_cell_to_node_matrix.rowConst(j);
          for (size_t R=0; R<cell_nodes.length; ++R) {
            int Rp1 = (R+1)%cell_nodes.length;
            int Rm1 = (R+cell_nodes.length-1)%cell_nodes.length;
            Rd half_xrp_xrm = 0.5*(xr(cell_nodes(Rp1))-xr(cell_nodes(Rm1)));
            m_Cjr(j,R) = Rd{-half_xrp_xrm[1], half_xrp_xrm[0]};
          }
        });

      const NodeValuePerCell<Rd>& Cjr = m_Cjr;
      Kokkos::parallel_for(m_Cjr.numberOfValues(), KOKKOS_LAMBDA(const int& j){
          m_ljr[j] = l2Norm(Cjr[j]);
        });

      const NodeValuePerCell<double>& ljr = m_ljr;
      Kokkos::parallel_for(m_Cjr.numberOfValues(), KOKKOS_LAMBDA(const int& j){
          m_njr[j] = (1./ljr[j])*Cjr[j];
        });

    } else if (dimension ==3) {
      const Kokkos::View<const Rd*> xr = m_mesh.xr();

#warning Rewrite using better data structures and remove this explicit 4
      Kokkos::View<Rd**> Nlr("Nlr", m_mesh.connectivity().numberOfFaces(), 4);

      Kokkos::parallel_for(m_mesh.numberOfFaces(), KOKKOS_LAMBDA(const int& l) {
          const auto& face_nodes = m_mesh.connectivity().m_face_to_node_matrix.rowConst(l);
          const size_t nb_nodes = face_nodes.length;
          std::vector<Rd> dxr(nb_nodes);
          for (size_t r=0; r<nb_nodes; ++r) {
            dxr[r] = xr[face_nodes((r+1)%nb_nodes)] - xr[face_nodes((r+nb_nodes-1)%nb_nodes)];
          }
          const double inv_12_nb_nodes = 1./(12.*nb_nodes);
          for (size_t r=0; r<nb_nodes; ++r) {
            Rd Nr = zero;
            const Rd two_dxr = 2*dxr[r];
            for (size_t s=0; s<nb_nodes; ++s) {
              Nr += crossProduct((two_dxr - dxr[s]), xr[face_nodes(s)]);
            }
            Nr *= inv_12_nb_nodes;
            Nr -= (1./6.)*crossProduct(dxr[r], xr[face_nodes(r)]);
            Nlr(l,r) = Nr;
          }
        });

      Kokkos::parallel_for(m_Cjr.numberOfValues(), KOKKOS_LAMBDA(const int& jr){
          m_Cjr[jr] = zero;
        });

      Kokkos::parallel_for(m_mesh.numberOfCells(), KOKKOS_LAMBDA(const int& j) {
          const auto& cell_nodes = m_mesh.connectivity().m_cell_to_node_matrix.rowConst(j);

          const auto& cell_faces = m_mesh.connectivity().m_cell_to_face_matrix.rowConst(j);
          const auto& cell_face_is_reversed = m_mesh.connectivity().m_cell_face_is_reversed.itemValues(j);

          for (size_t L=0; L<cell_faces.length; ++L) {
            const size_t l = cell_faces(L);
            const auto& face_nodes = m_mesh.connectivity().m_face_to_node_matrix.rowConst(l);

#warning should this lambda be replaced by a precomputed correspondance?
            std::function local_node_number_in_cell
                = [&](const size_t& node_number) {
                    for (size_t i_node=0; i_node<cell_nodes.length; ++i_node) {
                      if (node_number == cell_nodes(i_node)) {
                        return i_node;
                        break;
                      }
                    }
                    return std::numeric_limits<size_t>::max();
                  };

            if (cell_face_is_reversed[L]) {
              for (size_t rl = 0; rl<face_nodes.length; ++rl) {
                const size_t R = local_node_number_in_cell(face_nodes(rl));
                m_Cjr(j, R) -= Nlr(l,rl);
              }
            } else {
              for (size_t rl = 0; rl<face_nodes.length; ++rl) {
                const size_t R = local_node_number_in_cell(face_nodes(rl));
                m_Cjr(j, R) += Nlr(l,rl);
              }
            }
          }
        });

      const NodeValuePerCell<Rd>& Cjr = m_Cjr;
      Kokkos::parallel_for(m_Cjr.numberOfValues(), KOKKOS_LAMBDA(const int& jr){
          m_ljr[jr] = l2Norm(Cjr[jr]);
        });

      const NodeValuePerCell<double>& ljr = m_ljr;
      Kokkos::parallel_for(m_Cjr.numberOfValues(), KOKKOS_LAMBDA(const int& jr){
          m_njr[jr] = (1./ljr[jr])*Cjr[jr];
        });
    }
    static_assert((dimension<=3), "only 1d, 2d and 3d are implemented");
  }

 public:
  const MeshType& mesh() const
  {
    return m_mesh;
  }

  const NodeValuePerCell<Rd>& Cjr() const
  {
    return m_Cjr;
  }

  const NodeValuePerCell<double>& ljr() const
  {
    return m_ljr;
  }

  const NodeValuePerCell<Rd>& njr() const
  {
    return m_njr;
  }

  const Kokkos::View<const Rd*> xj() const
  {
    return m_xj;
  }

  const Kokkos::View<const double*> Vj() const
  {
    return m_Vj;
  }

  void updateAllData()
  {
    this->_updateCjr();
    this->_updateCenter();
    this->_updateVolume();
  }

  MeshData(const MeshType& mesh)
      : m_mesh(mesh),
        m_Cjr(mesh.connectivity()),
        m_ljr(mesh.connectivity()),
        m_njr(mesh.connectivity()),
        m_xj("xj", mesh.numberOfCells()),
        m_Vj("Vj", mesh.numberOfCells())
  {
    if constexpr (dimension==1) {
      // in 1d Cjr are computed once for all
      Kokkos::parallel_for(m_mesh.numberOfCells(), KOKKOS_LAMBDA(const int& j) {
          m_Cjr(j,0)=-1;
          m_Cjr(j,1)= 1;
        });
      // in 1d njr=Cjr (here performs a shallow copy)
      m_njr=m_Cjr;
      Kokkos::parallel_for(m_mesh.numberOfCells(), KOKKOS_LAMBDA(const int& j) {
          m_ljr(j,0)= 1;
          m_ljr(j,1)= 1;
        });
    }
    this->updateAllData();
  }

  ~MeshData()
  {
    ;
  }
};

#endif // MESH_DATA_HPP