#include "animation.hpp"

template <int D> class HardSphere : public Sphere<D> {
public:
  using Sphere<D>::Sphere;

  double U(const Position<D>& y, double ry) const {
    if (Sphere<D>::regularizedDistanceSquared(y, ry) < 1) {
      return 1e6;
    } else {
      return 0;
    }
  }
};

int main(int argc, char* argv[]) {
  unsigned N = 1200;
  double R = 0.024;
  double ε = 0.005;
  unsigned steps = 1e6;

  initializeAnimation(argc, argv);

  Rng r;

  Model<2, HardSphere<2>> m(N, 2 * R, gContinuousPolydisperse(R), r);

  for (unsigned i = 0; i < steps; i++) {
    for (unsigned j = 0; j < N; j++) {
      HardSphere<2>& p = r.pick(m.particles);
      Position<2> xNew = p.x;
      Vector<2> Δx;
      for (double& xi : Δx) {
        xi = r.variate<double, std::normal_distribution>(0, ε);
      }
      xNew += Δx;
      bool reject = false;
      for (const HardSphere<2>* neighbor : m.nl.neighborsOf(xNew)) {
        if (neighbor != &p) {
          if (neighbor->regularizedDistanceSquared(xNew, p.r) < 1 && neighbor->regularizedDistanceSquared(p.x, p.r) >= 1) {
            reject = true;
            break;
          }
        }
      }
      if (!reject) {
        m.move(p, xNew);
      }
      
      HardSphere<2>& p1 = r.pick(m.particles);
      HardSphere<2>& p2 = r.pick(m.particles);

      double rat = p1.r / p2.r;
      if (rat < 1.2 && rat > 0.83) {
        reject = false;
        for (const HardSphere<2>* neighbor : m.nl.neighborsOf(p2.x)) {
          if (neighbor != &p1 && neighbor != &p2) {
            if (neighbor->regularizedDistanceSquared(p2.x, p1.r) < 1 && neighbor->regularizedDistanceSquared(p1.x, p1.r) >= 1) {
              reject = true;
              break;
            }
          }
        }
        for (const HardSphere<2>* neighbor : m.nl.neighborsOf(p1.x)) {
          if (neighbor != &p1 && neighbor != &p2) {
            if (neighbor->regularizedDistanceSquared(p1.x, p2.r) < 1 && neighbor->regularizedDistanceSquared(p2.x, p2.r) >= 1) {
              reject = true;
              break;
            }
          }
        }
        if (!reject) {
          std::swap(p1.r, p2.r);
        }
      }
    }

    if (i % 10 == 0)
      draw(m);
  }

  return 0;
}