// Copyright (c) 2013-2020, SIB - Swiss Institute of Bioinformatics and
//                          Biozentrum - University of Basel
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// 
//   http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#ifndef PROMOD3_DYNAMIC_SPATIAL_ORGANIZER_HI
#define PROMOD3_DYNAMIC_SPATIAL_ORGANIZER_HI

#include <vector>
#include <map>

#include <ost/geom/geom.hh>

#include <promod3/core/message.hh>


namespace promod3 { namespace core {

template<class ITEM>
class SpatialOrganizerCell{

public:

  SpatialOrganizerCell() {
    x_.reserve(32);
    y_.reserve(32);
    z_.reserve(32);
    pointers_.reserve(32);
    distance_buffer_.reserve(32);
  }

  void AddItem(ITEM* item, Real x, Real y, Real z) {
    x_.push_back(x);
    y_.push_back(y);
    z_.push_back(z);
    pointers_.push_back(item);
    distance_buffer_.push_back(0.0);
  }

  void RemoveItem(ITEM* item) {
    typename std::vector<ITEM*>::iterator it = 
    std::find(pointers_.begin(), pointers_.end(), item);
    if(it == pointers_.end()) {
      throw promod3::Error("Cannot Remove inexistent item from spatial organizer cell!");
    }
    int idx = std::distance(pointers_.begin(), it);
    x_.erase(x_.begin() + idx);
    y_.erase(y_.begin() + idx);
    z_.erase(z_.begin() + idx);
    pointers_.erase(pointers_.begin() + idx);
    distance_buffer_.pop_back();
  }

  void ResetPos(ITEM* item, Real x, Real y, Real z) {
    typename std::vector<ITEM*>::iterator it = 
    std::find(pointers_.begin(), pointers_.end(), item);
    if(it == pointers_.end()) {
      throw promod3::Error("Cannot Remove inexistent item from spatial organizer cell!");
    }
    int idx = std::distance(pointers_.begin(), it);
    x_[idx] = x;
    y_[idx] = y;
    z_[idx] = z;
  }

  void FindClose(std::vector<std::pair<ITEM*,Real> >& result_vec,
                 Real dist, Real x, Real y, Real z) const{

    Real dx, dy, dz;
    Real squared_cutoff = dist*dist;
    int num_items = x_.size();

    for(int i = 0; i < num_items; ++i) {
      dx = x_[i] - x;
      dy = y_[i] - y;
      dz = z_[i] - z; 
      distance_buffer_[i] = dx*dx + dy*dy + dz*dz;
    }

    for(int i = 0; i < num_items; ++i) {
      if(distance_buffer_[i] <= squared_cutoff) {
        result_vec.push_back(std::make_pair(pointers_[i],std::sqrt(distance_buffer_[i])));
      }
    }
  }

private:
  std::vector<Real> x_;
  std::vector<Real> y_;
  std::vector<Real> z_;
  std::vector<ITEM*> pointers_;
  mutable std::vector<Real> distance_buffer_;
};

//! spatial organizer
/*
  organizes ITEMs defined by positions
  in a spatial hash map, into bins of
  a defined size (as given in the ctor)
*/
template <class ITEM>
class DynamicSpatialOrganizer {
public:
  typedef std::vector<std::pair<ITEM*,Real> > WithinList;
  typedef std::pair<std::pair<ITEM*,Real>*,uint> WithinResult;

private:
  struct Index {
    Index():
      u(0),v(0),w(0) {}
    Index(int uu, int vv, int ww): 
      u(uu),v(vv),w(ww) {}

    bool operator<(const Index& other) const {
      return w!=other.w ? w<other.w : (v!=other.v ? v<other.v : u<other.u);
    }

    bool operator==(const Index& other) const {
      return u == other.u && v == other.v && w == other.w;
    }

    bool operator!=(const Index& other) const {
      return u != other.u || v != other.v || w != other.w;
    }

    int u,v,w;

    static Index Max(const Index& a, const Index& b) {
      Index m;
      m.u = a.u > b.u ? a.u : b.u;
      m.v = a.v > b.v ? a.v : b.v;
      m.w = a.w > b.w ? a.w : b.w;
      return m;
    }

    static Index Min(const Index& a, const Index& b) {
      Index m;
      m.u = a.u < b.u ? a.u : b.u;
      m.v = a.v < b.v ? a.v : b.v;
      m.w = a.w < b.w ? a.w : b.w;
      return m;
    }
  };

  typedef std::map<Index,SpatialOrganizerCell<ITEM>*> ItemMap;


public:

  DynamicSpatialOrganizer(Real delta,
                          const std::vector<geom::Mat4>& transformations = 
                          std::vector<geom::Mat4>()): delta_(delta) {
    if(delta <= 0.0) {
      throw "delta must be positive!";
    }

    for(uint i = 0; i < transformations.size(); ++i) {
      geom::Mat4 inv_t = geom::Invert(transformations[i]);
      inverted_transformations_.push_back(inv_t);
    }
  }

  ~DynamicSpatialOrganizer() {
    for(typename ItemMap::iterator i = map_.begin(); i != map_.end(); ++i) {
      delete i->second;
    }
  }

  void Add(ITEM* item, const geom::Vec3& pos) {
    Index indx=gen_index(pos);
    typename ItemMap::iterator it = map_.find(indx);
    if(it == map_.end()) {
      map_[indx] = new SpatialOrganizerCell<ITEM>;
      it = map_.find(indx);
    }
    it->second->AddItem(item,pos[0],pos[1],pos[2]);
  }

  void Remove(ITEM* item, const geom::Vec3& pos) {
    Index indx=gen_index(pos);
    typename ItemMap::iterator i = map_.find(indx);
    if(i != map_.end()) {
      i->second->RemoveItem(item);
    }
    else{
      throw promod3::Error("Cannot remove non existent item from dynamic spatial organizer!");
    }
  }

  void Reset(ITEM* item, const geom::Vec3& actual_pos, const geom::Vec3& new_pos) {
    Index actual_indx=gen_index(actual_pos);
    Index new_indx=gen_index(new_pos);
    typename ItemMap::iterator i = map_.find(actual_indx);
    if(i == map_.end()) {
      throw promod3::Error("Could not find specified item in dynamic spatial organizer!"); 
    }
    if(actual_indx == new_indx) {
      i->second->ResetPos(item,new_pos[0],new_pos[1],new_pos[2]);
    } else {
      typename ItemMap::iterator j = map_.find(new_indx);
      if(j == map_.end()) {
        map_[new_indx] = new SpatialOrganizerCell<ITEM>;
        j = map_.find(new_indx);
      }
      i->second->RemoveItem(item);
      j->second->AddItem(item,new_pos[0],new_pos[1],new_pos[2]);
    }
  }

  WithinResult FindWithin(const geom::Vec3& pos, Real dist, 
                          bool clear_buffer = true) const{

    Index imin = gen_index(pos-geom::Vec3(dist,dist,dist));
    Index imax = gen_index(pos+geom::Vec3(dist,dist,dist));

    if(clear_buffer) {
      result_buffer_.clear();
    }

    int first_buffer_position = result_buffer_.size();

    for(int wc=imin.w;wc<=imax.w;++wc) {
      for(int vc=imin.v;vc<=imax.v;++vc) {
        for(int uc=imin.u;uc<=imax.u;++uc) {
          typename ItemMap::const_iterator map_it = map_.find(Index(uc,vc,wc));
          if(map_it!=map_.end()) {
            map_it->second->FindClose(result_buffer_, dist, pos[0], pos[1], pos[2]);
          }
        }
      }
    }

    WithinResult result;

    if(!result_buffer_.empty()) {
      result = std::make_pair(&result_buffer_[first_buffer_position],
                              result_buffer_.size() - first_buffer_position);
    }
    else{
      result = std::make_pair(static_cast<std::pair<ITEM*,Real>* >(NULL), 0);   
    }
    return result;
  }

  std::vector<WithinResult> FindWithinFull(const geom::Vec3& pos, 
                                           Real dist) {

    int n = inverted_transformations_.size();
    std::vector<WithinResult> return_vec(n + 1);    

    // do first the FindWithin on the actual non transformed positions
    return_vec[0] = this->FindWithin(pos, dist, true);

    for(int i = 0; i < n; ++i) {
      // instead of transforming all positions in the organizer, we apply
      // the corresponding inverted transform to the query position
      // Thanks David Sehnal for the trick!!!
      const geom::Mat4& t = inverted_transformations_[i];
      Real a = t(0,0)*pos[0] + t(0,1)*pos[1] + t(0,2)*pos[2] + t(0,3);
      Real b = t(1,0)*pos[0] + t(1,1)*pos[1] + t(1,2)*pos[2] + t(1,3);
      Real c = t(2,0)*pos[0] + t(2,1)*pos[1] + t(2,2)*pos[2] + t(2,3);      
      geom::Vec3 transformed_pos(a, b, c);
      return_vec[i + 1] = this->FindWithin(transformed_pos, dist, false);
    }

    return return_vec;
  }

  void Clear() {
    for(typename ItemMap::iterator i = map_.begin(); i != map_.end(); ++i) {
      delete i->second;
    }
    map_.clear();
  }

private:

  Real delta_;
  std::vector<geom::Mat4> inverted_transformations_;
  ItemMap map_;
  mutable WithinList result_buffer_;

  Index gen_index(const geom::Vec3& pos) const {
    Index nrvo(std::floor(pos[0]/delta_), std::floor(pos[1]/delta_),
               std::floor(pos[2]/delta_));
    return nrvo;
  }
};

}} // ns

#endif
