/**
 *  \file grid_close_pair_impl.h
 *  \brief Implementation of close pairs finder that is based on a grid hierarchy
 *
 *  Copyright 2007-2021 IMP Inventors. All rights reserved.
 *
 */

#ifndef IMPCORE_INTERNAL_GRID_CLOSE_PAIRS_IMPL_H
#define IMPCORE_INTERNAL_GRID_CLOSE_PAIRS_IMPL_H

#include <IMP/core/core_config.h>
#include "../GridClosePairsFinder.h"
#include "sinks.h"
#include "../QuadraticClosePairsFinder.h"
#include "../XYZR.h"
#include <IMP/algebra/standard_grids.h>
#include <IMP/core/utility.h>

IMPCORE_BEGIN_INTERNAL_NAMESPACE

//#define IMP_GRID_CPF_CHECKS

inline bool get_interiors_intersect(const algebra::Vector3D &v, double ra,
                                    double rb) {
  double sr = ra + rb;
  double a0 = std::abs(v[0]);
  if (a0 >= sr) return false;
  double a1 = std::abs(v[1]);
  if (a1 >= sr) return false;
  double a2 = std::abs(v[2]);
  if (a2 >= sr) return false;
  return v * v < square(sr);
}
inline bool get_interiors_intersect(const algebra::Sphere3D &s1,
                                    const algebra::Sphere3D &s2,
                                    double d) {
  double sr = s1.get_radius() + s2.get_radius() + d;
  algebra::Vector3D v(s1.get_center() - s2.get_center());
  // Probably not worth it all the CPU time on ifs
  /*  double a0 = std::abs(v[0]);
  if (a0 >= sr) return false;
  double a1 = std::abs(v[1]);
  if (a1 >= sr) return false;
  double a2 = std::abs(v[2]);
  if (a2 >= sr) return false;*/
  return v * v < square(sr);
}

struct ParticleTraits {
  typedef ParticleIndex ID;
  Model *m_;
  double d_;
  mutable algebra::Sphere3D const* model_spheres_table_; // faster access to spheres data - this assumes model spheres table remains the same within the lifetime of object

  ParticleTraits(Model *m, double d) : m_(m), d_(d),
    model_spheres_table_( m->access_spheres_data() )
  { }
  ParticleIndex get_id(Particle *p, int) const {
    return p->get_index();
  }
  algebra::Sphere3D const& get_sphere(ParticleIndex p) const{
    return model_spheres_table_[p.get_index()];
  }
  algebra::Vector3D get_center(ParticleIndex p, int) const {
    return get_sphere(p).get_center();
  }
  double get_radius(ParticleIndex p, int) const {
    return get_sphere(p).get_radius();
  }
  bool get_is_close(ParticleIndex a, ParticleIndex b) const {
    return get_interiors_intersect(get_sphere(a),
                                   get_sphere(b),
                                   d_);
  }
  bool check_close(ParticleIndex a, ParticleIndex b) const {
    return get_interiors_intersect(get_center(a, 0) - get_center(b, 0),
                                   .95 * (get_radius(b, 0) + d_),
                                   .95 * (get_radius(a, 0)));
  }
  double get_distance() const { return d_; }
};

struct ParticleIndexTraits : ParticleTraits {
  ParticleIndexTraits(Model *m, double d) : ParticleTraits(m, d) {}
  ParticleIndex get_id(ParticleIndex p, int) const { return p; }
};

struct BoundingBoxTraits {
  typedef unsigned int ID;
  algebra::BoundingBox3Ds::const_iterator it0_, it1_;
  double d_;
  BoundingBoxTraits(algebra::BoundingBox3Ds::const_iterator it0,
                    algebra::BoundingBox3Ds::const_iterator it1, double d)
      : it0_(it0), it1_(it1), d_(d) {}
  algebra::BoundingBox3Ds::const_iterator get_it(int which) const {
    switch (which) {
      case 0:
        return it0_;
      default:
        return it1_;
    }
  }
  unsigned int get_id(const algebra::BoundingBox3D &bb, int which) const {
    return &bb - &*get_it(which);
  }
  double get_radius(unsigned int i, int which) const {
    return algebra::get_maximum_length(*(get_it(which) + i)) / 2.0;
  }
  algebra::Vector3D get_center(unsigned int i, int which) const {
    return algebra::Vector3D(.5 * ((get_it(which) + i)->get_corner(0) +
                                   (get_it(which) + i)->get_corner(1)));
  }
  bool get_is_close(unsigned int a, unsigned int b) const {
    algebra::BoundingBox3D ag = *(it0_ + a) + d_;
    return algebra::get_interiors_intersect(ag, *(it1_ + b));
  }
  bool check_close(unsigned int, unsigned int) const { return false; }
  double get_distance() const { return d_; }
};

struct BBPairSink {
  IntPairs &out_;
  BBPairSink(IntPairs &out) : out_(out) {}
  typedef IntPair argument_type;
  bool operator()(unsigned int a, unsigned int b) const {
    out_.push_back(IntPair(a, b));
    return true;
  }
  bool check_contains(unsigned int, unsigned int) const { return true; }
};

template <class Traits>
struct Helper {
  typedef typename Traits::ID ID;
  //! an extension of vector, containing a (template) vector of
  //! Traits::ID, with an integer indentifier which_
  struct IDs : public Vector<ID> {
    int which_;
    //! an empty ids vector identified by which
    IDs(int which) : which_(which) {}
    //! an ids vector of one element with value id, identified by which
    IDs(ID id, int which) : Vector<ID>(1, id), which_(which) {}
    IDs() : which_(-1) {}
  };

  typedef typename algebra::SparseGrid3D<IDs> Grid;
  typedef Vector<Grid> Grids;

  template <class It>
  struct ParticleSet {
    It b_, e_;
    int which_;
    ParticleSet(It b, It e, int which) : b_(b), e_(e), which_(which) {}
    unsigned int size() const { return std::distance(b_, e_); }
  };
  template <class It>
  static ParticleSet<It> get_particle_set(It b, It e, int which) {
    return ParticleSet<It>(b, e, which);
  }
  template <class It>
  struct IDSet {
    It b_, e_;
    int which_;
    IDSet(It b, It e, int which) : b_(b), e_(e), which_(which) {}
  };
  template <class It>
  static IDSet<It> get_id_set(It b, It e, int which) {
    return IDSet<It>(b, e, which);
  }

  static int get_index(int ii, int which, int lb, int ub) {
    if (which == 1) {
      if (ii == lb)
        ii = ub;
      else if (ii == ub)
        ii = lb;
    }
    return ii;
  }

  //! returns the maximal radius of particles in ps0, using tr to compute
  //! that radius
  template <class It>
  static double get_max_radius(const ParticleSet<It> &ps0, const Traits &tr) {
    double maxr = 0;
    for (It c = ps0.b_; c != ps0.e_; ++c) {
      maxr =
          std::max(tr.get_radius(tr.get_id(*c, ps0.which_), ps0.which_), maxr);
    }
    return maxr;
  }

  static algebra::BoundingBox3D get_bb(const IDs &ids, const Traits &tr) {
    algebra::BoundingBox3D bb;
    for (typename IDs::const_iterator c = ids.begin(); c != ids.end(); ++c) {
      bb += tr.get_center(*c, ids.which_);
    }
    return bb;
  }

  template <class It>
  static void partition_points(const ParticleSet<It> &ps, const Traits &tr,
                               Vector<IDs> &bin_contents,
                               Floats &bin_ubs) {
    bin_contents.push_back(IDs(ps.which_));
    for (It c = ps.b_; c != ps.e_; ++c) {
      double cr = tr.get_radius(tr.get_id(*c, ps.which_), ps.which_) + 0;
      while (cr < std::max(.5 * bin_ubs.back() - .5 * tr.get_distance(), 0.0)) {
        double v = std::max(.5 * bin_ubs.back() - .5 * tr.get_distance(), 0.1);
        if (v > .1) {
          bin_ubs.push_back(v);
        } else
          break;
      }
      for (int i = bin_ubs.size() - 1; i >= 0; --i) {
        if (cr <= bin_ubs[i] || i == 0) {
          while (static_cast<int>(bin_contents.size()) < i + 1) {
            bin_contents.push_back(IDs(ps.which_));
          }
          bin_contents[i].push_back(tr.get_id(*c, ps.which_));
          break;
        }
      }
    }
    IMP_IF_CHECK(USAGE_AND_INTERNAL) {
      int total = 0;
      for (unsigned int i = 0; i < bin_contents.size(); ++i) {
        total += bin_contents[i].size();
      }
      IMP_CHECK_VARIABLE(total);
      IMP_INTERNAL_CHECK(total == std::distance(ps.b_, ps.e_),
                         "Lost points " << total << " "
                                        << std::distance(ps.b_, ps.e_));
    }
  }

  static Grid create_grid(const algebra::BoundingBox3D &bb, double side) {
    return Grid(side, bb);
  }

  //! ids - vector of IDs (e.g. Particle Indexes) identified by which_
  //! traits - traits class for finding close pairs of IDs
  //! g - a grid for storing the ids
  static void fill_grid(const IDs &ids, const Traits &tr, Grid &g) {
    for (typename IDs::const_iterator c = ids.begin(); c != ids.end(); ++c) {
      algebra::Vector3D v = tr.get_center(*c, ids.which_);
      typename Grid::ExtendedIndex ind = g.get_nearest_extended_index(v);
      if (g.get_has_index(ind)) {
        g[g.get_index(ind)].push_back(*c);
      } else {
        g.add_voxel(ind, IDs(*c, ids.which_));
      }
    }
    IMP_IF_LOG(VERBOSE) {
      IMP_LOG_VERBOSE("Grid built" << std::endl);
      for (typename Grid::AllConstIterator it = g.all_begin();
           it != g.all_end(); ++it) {
        IMP_INTERNAL_CHECK(it->second.size() > 0, "Empty voxel");
        IMP_LOG_VERBOSE("Voxel " << it->first << " has ");
        for (unsigned int i = 0; i < it->second.size(); ++i) {
          IMP_LOG_VERBOSE(Showable(it->second[i]) << " ");
        }
        IMP_LOG_VERBOSE(std::endl);
      }
    }
  }

  template <class It, class Out>
  static bool do_fill_close_pairs_from_list(It b, It e, const Traits &tr,
                                            Out out) {
    for (It c = b; c != e; ++c) {
      for (It cp = b; cp != c; ++cp) {
        if (tr.get_is_close(*c, *cp)) {
          /*IMP_LOG_VERBOSE( "Found pair " << Showable(*c) << " "
            << Showable(*cp) << std::endl);*/
          if (!out(*c, *cp)) {
            return false;
          }
        }
      }
    }
    return true;
  }

  template <class ItA, class ItB, class Out>
  static bool do_fill_close_pairs_from_lists(ItA ab, ItA ae, ItB bb, ItB be,
                                             const Traits &tr, Out out) {
    for (ItA c = ab; c != ae; ++c) {
      for (ItB cp = bb; cp != be; ++cp) {
        if (tr.get_is_close(*c, *cp)) {
          /*IMP_LOG_VERBOSE( "Found pair " << Showable(*c) << " "
            << Showable(*cp) << std::endl);*/
          if (!out(*c, *cp)) {
            return false;
          }
        }
      }
    }
    return true;
  }

  template <class Out>
  static bool do_fill_close_pairs(const Grid &gg, typename Grid::Index index,
                                  const IDs &qps, bool half, const Traits &tr,
                                  Out &out) {
    IMP::algebra::ExtendedGridIndexD<3> imp_min(std::max(0, index[0] - 1),
                                                std::max(0, index[1] - 1),
                                                std::max(0, index[2] - 1));
    IMP::algebra::ExtendedGridIndexD<3> imp_max(
        std::min(gg.get_number_of_voxels(0),
                 static_cast<unsigned int>(index[0] + 1)),
        std::min(gg.get_number_of_voxels(1),
                 static_cast<unsigned int>(index[1] + 1)),
        std::min(gg.get_number_of_voxels(2),
                 static_cast<unsigned int>(index[2] + 1)));
    IMP::algebra::ExtendedGridIndexD<3> imp_cur(0, 0, 0);
    if (half) {
      for (imp_cur[0] = imp_min[0]; imp_cur[0] <= index[0]; ++imp_cur[0]) {
        for (imp_cur[1] = imp_min[1]; imp_cur[1] <= imp_max[1]; ++imp_cur[1]) {
          if (imp_cur[0] == index[0] && imp_cur[1] > index[1]) break;
          for (imp_cur[2] = imp_min[2]; imp_cur[2] <= imp_max[2];
               ++imp_cur[2]) {
            if (imp_cur[0] == index[0] && imp_cur[1] == index[1] &&
                imp_cur[2] >= index[2])
              break;
            if (gg.get_has_index(imp_cur)) {
              IMP::algebra::GridIndexD<3> _1 = gg.get_index(imp_cur);
              IMP_LOG_VERBOSE("Checking pair "
                              << _1 << " " << index << ": "
                              << Showable(gg[_1]) << " and " << index
                              << " which is " << Showable(qps)
                              << std::endl);
              IMP_INTERNAL_CHECK(_1 != index, "Index returned by get nearby");
              if (!do_fill_close_pairs_from_lists(gg[_1].begin(), gg[_1].end(),
                                                  qps.begin(), qps.end(), tr,
                                                  out)) {
                return false;
              }
            }
          }
        }
      }
    } else {
      for (imp_cur[0] = imp_min[0]; imp_cur[0] <= imp_max[0]; ++imp_cur[0]) {
        for (imp_cur[1] = imp_min[1]; imp_cur[1] <= imp_max[1]; ++imp_cur[1]) {
          for (imp_cur[2] = imp_min[2]; imp_cur[2] <= imp_max[2];
               ++imp_cur[2]) {
            if (gg.get_has_index(imp_cur)) {
              IMP::algebra::GridIndexD<3> _1 = gg.get_index(imp_cur);
              IMP_LOG_VERBOSE("Checking pair "
                              << _1 << " " << index << ": "
                              << Showable(gg[_1]) << " and " << index
                              << " which is " << Showable(qps)
                              << std::endl);
              if (!do_fill_close_pairs_from_lists(gg[_1].begin(), gg[_1].end(),
                                                  qps.begin(), qps.end(), tr,
                                                  out)) {
                return false;
              }
            }
          }
        }
      }
    }
    if (half) {
      IMP_LOG_VERBOSE("Checking pair " << index << " " << index << std::endl);
      if (!do_fill_close_pairs_from_list(gg[index].begin(), gg[index].end(), tr,
                                         out)) {
        return false;
      }
    }
    return true;
  }

  template <class It, class Out>
  static bool fill_close_pairs(const ParticleSet<It> &ps, const Traits &tr,
                               Out out) {
    if (ps.size() == 0) return true;
    double maxr = get_max_radius(ps, tr);
    Vector<IDs> bin_contents_g;
    Floats bin_ubs;
    bin_ubs.push_back(maxr);

    partition_points(ps, tr, bin_contents_g, bin_ubs);

    IMP_LOG_VERBOSE("Divided points into " << bin_contents_g.size() << " bins ("
                                           << bin_ubs.size() << ") " << maxr
                                           << std::endl);
    IMP_IF_LOG(VERBOSE) {
      IMP_LOG_VERBOSE("For G, contents are " << std::endl);
      for (unsigned int i = 0; i < bin_contents_g.size(); ++i) {
        IMP_LOG_VERBOSE(i << " " << bin_ubs[i] << " "
                          << bin_contents_g[i].size() << std::endl);
      }
    }
    Vector<algebra::BoundingBox3D> bbs(bin_contents_g.size());
    for (unsigned int i = 0; i < bin_contents_g.size(); ++i) {
      bbs[i] = get_bb(bin_contents_g[i], tr);
    }
    for (unsigned int i = 0; i < bin_contents_g.size(); ++i) {
      if (bin_contents_g[i].empty()) continue;
      if (bin_contents_g[i].size() < 100) {
        if (!do_fill_close_pairs_from_list(bin_contents_g[i].begin(),
                                           bin_contents_g[i].end(), tr, out)) {
          return false;
        }
      } else {
        Grid gg = create_grid(
            bbs[i], std::max(tr.get_distance() + 2 * bin_ubs[i],
                             .01 * (bbs[i].get_corner(1) - bbs[i].get_corner(0))
                                       .get_magnitude()));
        fill_grid(bin_contents_g[i], tr, gg);
        for (typename Grid::AllConstIterator it = gg.all_begin();
             it != gg.all_end(); ++it) {
          if (!do_fill_close_pairs(gg, it->first, it->second, true, tr, out)) {
            return false;
          }
        }
      }
#ifdef IMP_GRID_CPF_CHECKS
      IMP_IF_CHECK(USAGE_AND_INTERNAL) {
        for (unsigned int k = 0; k < bin_contents_g[i].size(); ++k) {
          for (unsigned int j = 0; j < k; ++j) {
            if (tr.check_close(bin_contents_g[i][k], bin_contents_g[i][j])) {
              out.check_contains(bin_contents_g[i][k], bin_contents_g[i][j]);
            }
          }
        }
      }
#endif
      for (unsigned int j = 0; j < i; ++j) {
        if (bin_contents_g[j].empty()) continue;
        if (bin_contents_g[i].size() < 100 && bin_contents_g[j].size() < 100) {
          if (!do_fill_close_pairs_from_lists(
                   bin_contents_g[i].begin(), bin_contents_g[i].end(),
                   bin_contents_g[j].begin(), bin_contents_g[j].end(), tr,
                   out)) {
            return false;
          }
        } else {
          algebra::BoundingBox3D bb = bbs[i] + bbs[j];
          IMP_LOG_VERBOSE("Building grids for "
                          << i << " and " << j << " with bb " << bb
                          << " and side " << tr.get_distance() + bin_ubs[i] +
                                                 bin_ubs[j] << std::endl);
          Grid ggi(
              create_grid(bb, tr.get_distance() + bin_ubs[i] + bin_ubs[j]));
          Grid ggj = ggi;
          fill_grid(bin_contents_g[i], tr, ggi);
          fill_grid(bin_contents_g[j], tr, ggj);
          for (typename Grid::AllConstIterator it = ggj.all_begin();
               it != ggj.all_end(); ++it) {
            if (!do_fill_close_pairs(ggi, it->first, it->second, false, tr,
                                     out)) {
              return false;
            }
          }
        }
#ifdef IMP_GRID_CPF_CHECKS
        IMP_IF_CHECK(USAGE_AND_INTERNAL) {
          for (unsigned int k = 0; k < bin_contents_g[i].size(); ++k) {
            for (unsigned int l = 0; l < k; ++l) {
              if (tr.check_close(bin_contents_g[i][k], bin_contents_g[i][l])) {
                out.check_contains(bin_contents_g[i][k], bin_contents_g[i][l]);
              }
            }
          }
        }
#endif
      }
    }
#ifdef IMP_GRID_CPF_CHECKS
    IMP_IF_CHECK(USAGE_AND_INTERNAL) {
      for (It it0 = ps.b_; it0 != ps.e_; ++it0) {
        for (It it1 = ps.b_; it1 != it0; ++it1) {
          if (tr.check_close(tr.get_id(*it0, 0), tr.get_id(*it1, 1))) {
            out.check_contains(tr.get_id(*it0, 0), tr.get_id(*it1, 1));
          }
        }
      }
    }
#endif
    return true;
  }

  template <class ItG, class ItQ, class Out>
  static bool fill_close_pairs(const ParticleSet<ItG> &psg,
                               const ParticleSet<ItQ> &psq, const Traits &tr,
                               Out out) {
    if (psg.size() == 0 || psq.size() == 0) return true;
    double maxr = std::max(get_max_radius(psg, tr), get_max_radius(psq, tr));
    Vector<IDs> bin_contents_g, bin_contents_q;
    Floats bin_ubs;
    bin_ubs.push_back(maxr);

    partition_points(psg, tr, bin_contents_g, bin_ubs);
    partition_points(psq, tr, bin_contents_q, bin_ubs);
    IMP_LOG_VERBOSE("Divided points into " << bin_contents_g.size() << " and "
                                           << bin_contents_q.size() << " bins ("
                                           << bin_ubs.size() << ") " << maxr
                                           << std::endl);
    IMP_IF_LOG(VERBOSE) {
      IMP_LOG_VERBOSE("For G, contents are " << std::endl);
      for (unsigned int i = 0; i < bin_contents_g.size(); ++i) {
        IMP_LOG_VERBOSE(i << " " << bin_ubs[i] << " "
                          << bin_contents_g[i].size() << std::endl);
      }
      IMP_LOG_VERBOSE("For Q, contents are " << std::endl);
      for (unsigned int i = 0; i < bin_contents_q.size(); ++i) {
        IMP_LOG_VERBOSE(i << " " << bin_ubs[i] << " "
                          << bin_contents_q[i].size() << std::endl);
      }
    }
    Vector<algebra::BoundingBox3D> bbs_g(bin_contents_g.size());
    for (unsigned int i = 0; i < bin_contents_g.size(); ++i) {
      bbs_g[i] = get_bb(bin_contents_g[i], tr);
    }
    Vector<algebra::BoundingBox3D> bbs_q(bin_contents_q.size());
    for (unsigned int i = 0; i < bin_contents_q.size(); ++i) {
      bbs_q[i] = get_bb(bin_contents_q[i], tr);
    }
    for (unsigned int i = 0; i < bin_contents_g.size(); ++i) {
      if (bin_contents_g[i].empty()) continue;
      for (unsigned int j = 0; j < bin_contents_q.size(); ++j) {
        if (bin_contents_q[j].empty()) continue;
        algebra::BoundingBox3D bb = bbs_g[i] + bbs_q[j];
        if (bin_contents_g[i].size() < 10 && bin_contents_q[j].size() < 10) {
          if (!do_fill_close_pairs_from_lists(
                   bin_contents_g[i].begin(), bin_contents_g[i].end(),
                   bin_contents_q[j].begin(), bin_contents_q[j].end(), tr,
                   out)) {
            return false;
          }
        } else {
          IMP_LOG_VERBOSE("Building grids for "
                          << i << " and " << j << " with bb " << bb
                          << " and side " << tr.get_distance() + bin_ubs[i] +
                                                 bin_ubs[j] << std::endl);
          Grid gg =
              create_grid(bb, tr.get_distance() + bin_ubs[i] + bin_ubs[j]);
          Grid gq = gg;
          fill_grid(bin_contents_g[i], tr, gg);
          fill_grid(bin_contents_q[j], tr, gq);
          IMP_IF_CHECK(USAGE) {
            for (unsigned int ic = 0; ic < 3; ++ic) {
              IMP_USAGE_CHECK(
                  gg.get_number_of_voxels(ic) == gq.get_number_of_voxels(ic),
                  "Do not match on dimension " << ic << ": "
                  << gg.get_number_of_voxels(ic) << " vs "
                  << gq.get_number_of_voxels(ic));
            }
          }
          for (typename Grid::AllConstIterator it = gq.all_begin();
               it != gq.all_end(); ++it) {
            if (!do_fill_close_pairs(gg, it->first, it->second, false, tr,
                                     out)) {
              return false;
            }
          }
        }
#ifdef IMP_GRID_CPF_CHECKS
        IMP_IF_CHECK(USAGE_AND_INTERNAL) {
          for (unsigned int k = 0; k < bin_contents_g[i].size(); ++k) {
            for (unsigned int l = 0; l < bin_contents_q[j].size(); ++l) {
              if (tr.check_close(bin_contents_g[i][k], bin_contents_q[j][l])) {
                out.check_contains(bin_contents_g[i][k], bin_contents_q[j][l]);
              }
            }
          }
        }
#endif
      }
    }
#ifdef IMP_GRID_CPF_CHECKS
    IMP_IF_CHECK(USAGE_AND_INTERNAL) {
      for (ItG it0 = psg.b_; it0 != psg.e_; ++it0) {
        for (ItQ it1 = psq.b_; it1 != psq.e_; ++it1) {
          if (tr.check_close(tr.get_id(*it0, 0), tr.get_id(*it1, 1))) {
            out.check_contains(tr.get_id(*it0, 0), tr.get_id(*it1, 1));
          }
        }
      }
    }
#endif
    return true;
  }
};

typedef Helper<ParticleTraits> ParticleHelper;
typedef Helper<ParticleIndexTraits> ParticleIndexHelper;
typedef Helper<BoundingBoxTraits> BBHelper;

IMPCORE_END_INTERNAL_NAMESPACE

#endif /* IMPCORE_INTERNAL_GRID_CLOSE_PAIRS_IMPL_H */
