/**
 *  \file rigid_pair_score.h
 *  \brief utilities for rigid pair scores.
 *
 *  Copyright 2007-2021 IMP Inventors. All rights reserved.
 */

#ifndef IMPCORE_INTERNAL_RIGID_BODY_TREE_H
#define IMPCORE_INTERNAL_RIGID_BODY_TREE_H

#include <IMP/core/core_config.h>
#include "../XYZ.h"
#include "../rigid_bodies.h"
#include <IMP/algebra/Sphere3D.h>
#include <queue>

IMPCORE_BEGIN_INTERNAL_NAMESPACE

class IMPCOREEXPORT RigidBodyHierarchy : public IMP::Object {
  RigidBody rb_;
  struct Data {
    Ints children_;
    algebra::Sphere3D s_;
  };
  Vector<Data> tree_;
  ParticleIndexes constituents_;

  typedef Vector<unsigned int> SphereIndexes;
  typedef Vector<SphereIndexes> SpheresSplit;
  SpheresSplit divide_spheres(const algebra::Sphere3Ds &ss,
                              const SphereIndexes &s);
  void set_sphere(unsigned int ni, const algebra::Sphere3D &s);
  void set_leaf(unsigned int ni, const ParticleIndexes &ids);
  unsigned int add_children(unsigned int ni, unsigned int num_children);
  ParticleIndexes validate_internal(Model *m, int cur,
                                            algebra::Sphere3Ds bounds) const;
  typedef std::pair<unsigned int, SphereIndexes> Node;
  void add_node_as_leaf(Model *m, const Node &cur);
  void build_tree(Model *m, const Node &n,
                  const algebra::Sphere3Ds &spheres, Vector<Node> &stack);

 public:
  algebra::Sphere3D get_sphere(unsigned int i) const {
    IMP_INTERNAL_CHECK(i < tree_.size(), "Out of spheres vector");
    IMP_CHECK_OBJECT(rb_.get_particle());
    algebra::Sphere3D ret(rb_.get_reference_frame().get_global_coordinates(
                              tree_[i].s_.get_center()),
                          tree_[i].s_.get_radius());
    return ret;
  }
  bool get_is_leaf(unsigned int ni) const {
    IMP_INTERNAL_CHECK(ni < tree_.size(), "Out of range");
    IMP_INTERNAL_CHECK(!tree_[ni].children_.empty(),
                       "Everything must have particles or children");
    return tree_[ni].children_[0] < 0;
  }
  unsigned int get_number_of_particles(unsigned int ni) const {
    IMP_INTERNAL_CHECK(ni < tree_.size(), "Out of range");
    IMP_INTERNAL_CHECK(get_is_leaf(ni), "Only leaves have particles");
    return tree_[ni].children_.size();
  }
  unsigned int get_number_of_children(unsigned int ni) const {
    IMP_INTERNAL_CHECK(ni < tree_.size(), "Out of range");
    if (!get_is_leaf(ni)) {
      return tree_[ni].children_.size();
    } else {
      return 1;
    }
  }
  unsigned int get_child(unsigned int ni, unsigned int i) const {
    IMP_INTERNAL_CHECK(ni < tree_.size(), "Out of range");
    IMP_INTERNAL_CHECK(tree_[ni].children_.size() > i,
                       "Out of range in particle");
    if (!get_is_leaf(ni)) {
      return tree_[ni].children_[i];
    } else {
      return ni;
    }
  }
  ParticleIndex get_particle(unsigned int ni, unsigned int i) const {
    IMP_INTERNAL_CHECK(ni < tree_.size(), "Out of range");
    IMP_INTERNAL_CHECK(tree_[ni].children_.size() > i,
                       "Out of range in particle");
    IMP_INTERNAL_CHECK(tree_[ni].children_[i] < 0, "Not a leaf node");
    ParticleIndex index(std::abs(tree_[ni].children_[i]) - 1);
    return index;
  }
  algebra::Sphere3Ds get_all_spheres() const;
  RigidBodyHierarchy(RigidBody rb, const ParticleIndexes &constituents);
  algebra::Sphere3Ds get_tree() const;
  bool get_constituents_match(const ParticleIndexes &ps) const {
    if (ps.size() != constituents_.size()) return false;
    ParticleIndexes un;
    std::set_union(ps.begin(), ps.end(), constituents_.begin(),
                   constituents_.end(), std::back_inserter(un));
    return (un.size() == ps.size());
  }
  const ParticleIndexes &get_constituents() const {
    return constituents_;
  }
  // for testing
  ParticleIndexes get_particles(unsigned int pi) const {
    ParticleIndexes ret;
    Ints stack(1, pi);
    do {
      unsigned int i = stack.back();
      stack.pop_back();
      if (get_is_leaf(i)) {
        for (unsigned int j = 0; j < get_number_of_particles(i); ++j) {
          ret.push_back(get_particle(i, j));
        }
      } else {
        for (unsigned int j = 0; j < get_number_of_children(i); ++j) {
          stack.push_back(get_child(i, j));
        }
      }
    } while (!stack.empty());
    return ret;
  }
  void validate(Model *m) const;
  RigidBody get_rigid_body() const { return rb_; }
  Model *get_model() const { return rb_->get_model(); }
  IMP_OBJECT_METHODS(RigidBodyHierarchy);
};

IMPCOREEXPORT Particle *closest_particle(Model *m,
                                                 const RigidBodyHierarchy *da,
                                                 XYZR pt);

IMPCOREEXPORT Particle *closest_particle(
    Model *m, const RigidBodyHierarchy *da, XYZR pt,
    double dist = std::numeric_limits<double>::max());

IMPCOREEXPORT ParticlePair closest_pair(Model *m,
                                                const RigidBodyHierarchy *da,
                                                const RigidBodyHierarchy *db);
IMPCOREEXPORT ParticlePair closest_pair(
    Model *m, const RigidBodyHierarchy *da,
    const RigidBodyHierarchy *db,
    double dist = std::numeric_limits<double>::max());

struct LessFirst {
  template <class A>
  bool operator()(const A &a, const A &b) const {
    return a.first < b.first;
  }
};
inline double distance_bound(Model *m, const RigidBodyHierarchy *da,
                             unsigned int i, ParticleIndex b) {
  algebra::Sphere3D s = da->get_sphere(i);
  double rd = algebra::get_distance(s, m->get_sphere(b));
  return rd;
}

inline double distance_bound(Model *, const RigidBodyHierarchy *da,
                             unsigned int i, const RigidBodyHierarchy *db,
                             unsigned int j) {
  algebra::Sphere3D sa = da->get_sphere(i);
  algebra::Sphere3D sb = db->get_sphere(j);
  double rd = algebra::get_distance(sa, sb);
  return rd;
}

template <class Sink>
inline void fill_close_pairs(Model *m, const RigidBodyHierarchy *da,
                             const RigidBodyHierarchy *db, double dist,
                             Sink sink) {
  IMP_IF_CHECK(USAGE_AND_INTERNAL) {
    da->validate(m);
    db->validate(m);
  }
  typedef std::pair<int, int> IP;
  typedef std::pair<double, IP> QP;
  std::priority_queue<QP, Vector<QP>, LessFirst> queue;
  {
    double d = distance_bound(m, da, 0, db, 0);
    if (d < dist) {
      queue.push(QP(d, IP(0, 0)));
    }
  }
  while (!queue.empty()) {
    QP v = queue.top();
    queue.pop();
    /*IMP_LOG_TERSE( "Trying pair " << v.second.first << " " << v.second.second
      << std::endl);*/
    if (da->get_is_leaf(v.second.first) && db->get_is_leaf(v.second.second)) {
      for (unsigned int i = 0; i < da->get_number_of_particles(v.second.first);
           ++i) {
        ParticleIndex deca(da->get_particle(v.second.first, i));
        for (unsigned int j = 0;
             j < db->get_number_of_particles(v.second.second); ++j) {
          ParticleIndex decb(db->get_particle(v.second.second, j));
          double d =
              algebra::get_distance(m->get_sphere(deca), m->get_sphere(decb));
          if (d < dist) {
            if (!sink(deca, decb)) {
              return;
            }
          }
        }
      }
    } else if (da->get_is_leaf(v.second.first)) {
      for (unsigned int j = 0; j < db->get_number_of_children(v.second.second);
           ++j) {
        unsigned int child = db->get_child(v.second.second, j);
        double d = distance_bound(m, da, v.second.first, db, child);
        if (d < dist) {
          queue.push(QP(d, IP(v.second.first, child)));
        }
      }
    } else if (db->get_is_leaf(v.second.second)) {
      for (unsigned int i = 0; i < da->get_number_of_children(v.second.first);
           ++i) {
        unsigned int child = da->get_child(v.second.first, i);
        double d = distance_bound(m, da, child, db, v.second.second);
        if (d < dist) {
          queue.push(QP(d, IP(child, v.second.second)));
        }
      }
    } else {
      for (unsigned int i = 0; i < da->get_number_of_children(v.second.first);
           ++i) {
        unsigned int childa = da->get_child(v.second.first, i);
        for (unsigned int j = 0;
             j < db->get_number_of_children(v.second.second); ++j) {
          unsigned int childb = db->get_child(v.second.second, j);
          double d = distance_bound(m, da, childa, db, childb);
          if (d < dist) {
            queue.push(QP(d, IP(childa, childb)));
          }
        }
      }
    }
  }

  IMP_IF_CHECK(USAGE_AND_INTERNAL) {
    if (da->get_constituents().size() * db->get_constituents().size() < 1000) {
      for (unsigned int i = 0; i < da->get_constituents().size(); ++i) {
        XYZR ca(m, da->get_constituents()[i]);
        for (unsigned int j = 0; j < db->get_constituents().size(); ++j) {
          XYZR cb(m, db->get_constituents()[j]);
          if (get_distance(ca, cb) < .9 * dist) {
            sink.check_contains(da->get_constituents()[i],
                                db->get_constituents()[j]);
          }
        }
      }
    }
  }
}

IMPCOREEXPORT ParticlePairsTemp close_pairs(
    Model *m, const RigidBodyHierarchy *da,
    const RigidBodyHierarchy *db, double dist);

template <class Sink>
inline void fill_close_particles(Model *m, const RigidBodyHierarchy *da,
                                 ParticleIndex pt, double dist,
                                 Sink sink) {
  typedef std::pair<double, int> QP;
  std::priority_queue<QP, Vector<QP>, LessFirst> queue;
  {
    double d = distance_bound(m, da, 0, pt);
    if (d > dist) return;
    queue.push(QP(d, 0));
  }
  do {
    std::pair<double, int> v = queue.top();
    queue.pop();
    if (da->get_is_leaf(v.second)) {
      for (unsigned int i = 0; i < da->get_number_of_particles(v.second); ++i) {
        ParticleIndex p = da->get_particle(v.second, i);
        double d = algebra::get_distance(m->get_sphere(p), m->get_sphere(pt));
        if (d < dist) {
          if (!sink(p)) {
            return;
          }
        }
      }
    } else {
      for (unsigned int i = 0; i < da->get_number_of_children(v.second); ++i) {
        unsigned int c = da->get_child(v.second, i);
        double d = distance_bound(m, da, c, pt);
        if (d < dist) {
          queue.push(QP(d, c));
        }
      }
    }
  } while (!queue.empty());

  IMP_IF_CHECK(USAGE_AND_INTERNAL) {
    for (unsigned int i = 0; i < da->get_constituents().size(); ++i) {
      XYZR c(m, da->get_constituents()[i]);
      if (get_distance(c, XYZR(m, pt)) < .9 * dist) {
        sink.check_contains(da->get_constituents()[i]);
      }
    }
  }
}

IMPCOREEXPORT ParticlesTemp close_particles(
    Model *m, const RigidBodyHierarchy *da, XYZR pt, double dist);

IMPCOREEXPORT RigidBodyHierarchy *get_rigid_body_hierarchy(
    RigidBody rb, ParticleIndexes constituents,
    ObjectKey mykey = ObjectKey());

IMPCORE_END_INTERNAL_NAMESPACE

#endif /* IMPCORE_INTERNAL_RIGID_BODY_TREE_H */
