#include <sctl.hpp>

namespace SCTL_NAMESPACE {

template <class Real, Integer DIM> class ParticleTreeWrapper {
  public:

    ParticleTreeWrapper(Comm comm = Comm::World()) : tree_(comm) {}

    void GetTree(const Real*& node_coord, const int32_t*& parent_lst, const int32_t*& child_lst, const int8_t*& node_depth, const int8_t*& node_ghost, const int8_t*& node_leaf, int32_t& Nnodes) const {
      static constexpr Integer MAX_CHILD = (1u << DIM);
      const auto& node_mid_  = tree_.GetNodeMID();
      const auto& node_attr_ = tree_.GetNodeAttr();
      SCTL_ASSERT(node_mid_.Dim() == node_attr_.Dim());
      Nnodes = node_mid_.Dim();

      if (node_coord_.Dim() != Nnodes*DIM      ) node_coord_.ReInit(Nnodes*DIM      );
      if (parent_lst_.Dim() != Nnodes          ) parent_lst_.ReInit(Nnodes          );
      if (child_lst_ .Dim() != Nnodes*MAX_CHILD) child_lst_ .ReInit(Nnodes*MAX_CHILD);
      if (node_depth_.Dim() != Nnodes          ) node_depth_.ReInit(Nnodes          );
      if (node_ghost_.Dim() != Nnodes          ) node_ghost_.ReInit(Nnodes          );
      if (node_leaf_ .Dim() != Nnodes          ) node_leaf_ .ReInit(Nnodes          );
      Vector<Long> ancestors(Morton<DIM>::MAX_DEPTH);
      Vector<Long> child_cnt(Morton<DIM>::MAX_DEPTH);
      #pragma omp parallel for schedule(static)
      for (Long i = 0; i < Nnodes; i++) {
        node_mid_[i].template Coord<Real>(node_coord_.begin() + i*DIM);
        node_depth_[i] = node_mid_[i].Depth();
        node_ghost_[i] = node_attr_[i].Ghost;
        node_leaf_ [i] = node_attr_[i].Leaf;
        parent_lst_[i] = -1;
        for (Integer j = 0; j < MAX_CHILD; j++) child_lst_[i*MAX_CHILD+j] = -1;
      }
      for (Long i = 0; i < Nnodes; i++) { // Set parent_lst, child_lst_
        Integer depth = node_mid_[i].Depth();
        ancestors[depth] = i;
        child_cnt[depth] = 0;
        if (depth) {
          Long p = ancestors[depth-1];
          Long& c = child_cnt[depth-1];
          parent_lst_[i] = p+1; // +1 because indexing starts from 1 in fortran
          child_lst_[p*MAX_CHILD+c] = i+1; // +1 because indexing starts from 1 in fortran
          c++;
        }
      }

      node_coord = &node_coord_[0];
      parent_lst = &parent_lst_[0];
      child_lst  = &child_lst_ [0];
      node_depth = &node_depth_[0];
      node_ghost = &node_ghost_[0];
      node_leaf  = &node_leaf_ [0];
    }

    void UpdateRefinement(const Real* pt_coord_, int32_t Npt, int32_t max_pts, bool level_restrict, bool periodic) {
      const Vector<Real> pt_coord(Npt*DIM, Ptr2Itr<Real>((Real*)pt_coord_,Npt*DIM), false);
      tree_.UpdateRefinement(pt_coord, max_pts, level_restrict, periodic);
    }

    void AddData(const char* name, const Real* data, int32_t Ndata, const int32_t* cnt, int32_t Ncnt) {
      SCTL_ASSERT(Ncnt == tree_.GetNodeMID().Dim());
      const Vector<Real> data_(Ndata, Ptr2Itr<Real>((Real*)data, Ndata), false);
      Vector<Long> cnt_(Ncnt);
      for (Long i = 0; i < Ncnt; i++) cnt_[i] = cnt[i];
      tree_.AddData(name, data_, cnt_);
    }

    void GetData(Real*& data, int32_t& Ndata, const int32_t*& cnt, int32_t& Ncnt, const char* name) {
      Vector<Real> data_;
      Vector<Long> cnt_;
      tree_.GetData(data_, cnt_, name);
      data = &data_[0];
      Ndata = data_.Dim();
      Ncnt = cnt_.Dim();

      Vector<int32_t>& data_cnt = data_cnt_[name];
      data_cnt.ReInit(Ncnt);
      for (Long i = 0; i < Ncnt; i++) data_cnt[i] = cnt_[i];
      cnt = &data_cnt[0];
    }

    void DeleteData(const char* name) {
      tree_.DeleteData(name);
    }

    void WriteTreeVTK(const char* fname, bool show_ghost) const {
      tree_.WriteTreeVTK(fname, show_ghost);
    }


    void AddParticles(const char* name, const Real* coord, int32_t Npt) {
      const Vector<Real> coord_(Npt*DIM, Ptr2Itr<Real>((Real*)coord,Npt*DIM), false);
      tree_.AddParticles(name, coord_);
    }

    void AddParticleData(const char* name, const char* particle_name, const Real* data, int32_t Ndata) {
      const Vector<Real> data_(Ndata, Ptr2Itr<Real>((Real*)data,Ndata), false);
      tree_.AddParticleData(name, particle_name, data_);
    }

    void GetParticleData(Real*& data, int32_t& N, const char* name) const {
      Vector<Real>& data_ = pt_data_[name];
      tree_.GetParticleData(data_, name);
      N = data_.Dim();
      data = &data_[0];
    }

    void DeleteParticleData(const char* name) {
      tree_.DeleteParticleData(name);
    }

    void WriteParticleVTK(const char* fname, const char* data_name, bool show_ghost) const {
      tree_.WriteParticleVTK(fname, data_name, show_ghost);
    }

  private:
    PtTree<Real,DIM> tree_;

    mutable Vector<Real> node_coord_;
    mutable Vector<int32_t> parent_lst_;
    mutable Vector<int32_t> child_lst_;
    mutable Vector<int8_t> node_depth_;
    mutable Vector<int8_t> node_ghost_;
    mutable Vector<int8_t> node_leaf_;

    mutable std::map<std::string,Vector<int32_t>> data_cnt_;
    mutable std::map<std::string,Vector<Real>> pt_data_;
};

}

#ifdef __cplusplus
extern "C" {
#endif

void createtree_(void** tree) {
  tree[0] = (void*) new sctl::ParticleTreeWrapper<double,3>(sctl::Comm::World());
}

void deletetree_(void** tree) {
  delete (sctl::ParticleTreeWrapper<double,3>*)tree[0];
  tree[0] = nullptr;
}

void gettree_(const double** node_coord, const int32_t** parent_lst, const int32_t** child_lst, const int8_t** node_depth, const int8_t** node_ghost, const int8_t** node_leaf, int32_t* Nnodes, const void** tree_ctx) {
  const sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.GetTree(*node_coord, *parent_lst, *child_lst, *node_depth, *node_ghost, *node_leaf, *Nnodes);
}

void updaterefinement_(const double* pt_coord_, const int32_t* Npt, const int32_t* max_pts, bool* level_restrict, bool* periodic, void** tree_ctx) {
  sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.UpdateRefinement(pt_coord_, *Npt, *max_pts, *level_restrict, *periodic);
}

void adddata_(const char* name, const double* data, const int32_t* Ndata, const int32_t* cnt, const int32_t* Ncnt, void** tree_ctx) {
  sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.AddData(name, data, *Ndata, cnt, *Ncnt);
}

void getdata_(double** data, int32_t* Ndata, const int32_t** cnt, int32_t* Ncnt, const char* name, void** tree_ctx) {
  sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.GetData(*data, *Ndata, *cnt, *Ncnt, name);
}

void deletedata_(const char* name, void** tree_ctx) {
  sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.DeleteData(name);
}

void writetreevtk_(const char* fname, const bool* show_ghost, const void** tree_ctx) {
  const sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.WriteTreeVTK(fname, *show_ghost);
}


void addparticles_(const char* name, const double* coord, const int32_t* Npt, void** tree_ctx) {
  sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.AddParticles(name, coord, *Npt);
}

void addparticledata_(const char* name, const char* particle_name, const double* data, const int32_t* Ndata, void** tree_ctx) {
  sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.AddParticleData(name, particle_name, data, *Ndata);
}

void getparticledata_(double** data, int32_t* N, const char* name, const void** tree_ctx) {
  const sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.GetParticleData(*data, *N, name);
}

void deleteparticledata_(const char* name, void** tree_ctx) {
  sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.DeleteParticleData(name);
}

void writeparticlevtk_(const char* fname, const char* data_name, bool* show_ghost, const void** tree_ctx) {
  const sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
  tree.WriteParticleVTK(fname, data_name, *show_ghost);
}

#ifdef __cplusplus
}
#endif