#pragma once

#include <assert.h>
#include <set>
#include <vector>
#include <queue>

#include <eigen3/Eigen/Dense>

#include "pcg-cpp/include/pcg_random.hpp"
#include "randutils/randutils.hpp"

using Rng = randutils::random_generator<pcg32>;

template <int D> using Vector = Eigen::Matrix<double, D, 1>;
template <int D> using Matrix = Eigen::Matrix<double, D, D>;

template <int D> class Position : public Vector<D> {
private:
  void recenter() {
    for (double& vi : *this) {
      if (vi >= 0) {
        vi = fmod(vi, 1);
      } else {
        vi = fmod(vi + ceil(-vi), 1);
      }
    }
  }

public:
  using Vector<D>::Vector;

  Position<D>& operator+=(const Vector<D>& u) {
    Vector<D>::operator+=(u);
    recenter();
    return *this;
  }

  friend Position<D> operator+(Position<D> v, const Vector<D>& u) {
    v += u;
    return v;
  }

  Position<D>& operator-=(const Vector<D>& u) {
    Vector<D>::operator-=(u);
    recenter();
    return *this;
  }

  friend Vector<D> operator-(Position<D> v, const Position<D>& u) {
    v -= (Vector<D>)u;

    for (double& vi : v) {
      if (vi > 0.5)
        vi = 1 - vi;
    }

    return (Vector<D>)v;
  }
};

template <typename T, int D> concept Particle = requires(T v, const T& w) {
  { v.x } -> std::same_as<Position<D>>;
  { v.m } -> std::same_as<bool>;
  { v.U(w) } -> std::same_as<double>;
};

template <int D, Particle<D> T> class NeighborLists {
private:
  unsigned n;
  std::vector<std::set<T*>> lists;

  unsigned mod(signed a, unsigned b) const {
    if (a >= 0) {
      return a % b;
    } else {
      return (a + b * ((b - a) / b)) % b;
    }
  }

  int iPow(int x, unsigned p) const {
    if (p == 0)
      return 1;
    if (p == 1)
      return x;

    int tmp = iPow(x, p / 2);
    if (p % 2 == 0) {
      return tmp * tmp;
    } else {
      return x * tmp * tmp;
    }
  }

  unsigned index(Position<D> x) const {
    unsigned ind = 0;
    for (unsigned i = 0; i < D; i++) {
      ind += pow(n, i) * std::floor(n * x(i));
    }
    assert(ind < lists.size());
    return ind;
  }

  std::pair<typename std::set<T*>::iterator, bool> insert(T& p, unsigned ind) {
    return lists[ind].insert(&p);
  }

  size_t erase(T& p, unsigned ind) {
    return lists[ind].erase(&p);
  }

  void neighborsOfHelper(signed i0, unsigned d, std::set<T*>& ns) const {
    if (d == 0) {
      const std::set<T*>& a = lists[i0];
      ns.insert(a.begin(), a.end());
    } else {
      for (signed j : {-1, 0, 1}) {
        signed i1 = iPow(n, d) * (i0 / iPow(n, d)) + mod(i0 + j * iPow(n, d-1), iPow(n, d));
        neighborsOfHelper(i1, d - 1, ns);
      }
    }
  }

public:
  NeighborLists(unsigned m) : n(m), lists(pow(n, D)) {}
  NeighborLists(double l) : NeighborLists((unsigned)floor(1 / l)) {}

  std::pair<typename std::set<T*>::iterator, bool> insert(T& p, const Position<D>& x) {
    return insert(p, index(x));
  }

  std::pair<typename std::set<T*>::iterator, bool> insert(T& p) {
    return insert(p, p.x);
  }

  size_t erase(T& p, const Position<D>& x) {
    return erase(p, index(x));
  }

  size_t erase(T& p) {
    return erase(p, p.x);
  }

  void move(T& p, const Position<D>& xNew) {
    unsigned iOld = index(p.x);
    unsigned iNew = index(xNew);
    if (iNew != iOld) {
      erase(p, iOld);
      insert(p, iNew);
    }
  }

  std::set<T*> neighborsOf(const Position<D>& x) const {
    std::set<T*> neighbors;
    neighborsOfHelper(index(x), D, neighbors);
    return neighbors;
  }
};

template <int D> class Sphere {
public:
  Position<D> x;
  bool m;
  double r;

  Sphere(const Position<D>& x, double r) : x(x), m(false), r(r) {};

  double regularizedDistanceSquared(const Sphere<D>& s) const {
    return (x - s.x).squaredNorm() / pow(r + s.r, 2);
  } 

  bool intersectsBoundary() const {
    for (double xi : x) {
      if (xi < r || 1 - xi < r) {
        return true;
      }
    }
    return false;
  }
};

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

  double U(const HardSphere<D>& s) const {
    if (Sphere<D>::regularizedDistanceSquared(s) < 1) {
      return 1e6;
    } else {
      return 0;
    }
  }
};

template <int D, int n> class SoftSphere : public Sphere<D> {
private:
  const double c0 = -pow(2, 2*n-3) * pow(5, -n) * (8 + 6*n + pow(n, 2));
  const double c2 =  pow(2, 2+2*n) * pow(5, -n-2) * n * (4+n);
  const double c4 = -pow(2, 5+2*n) * pow(5, -n-4) * n * (2+n);
public:
  using Sphere<D>::Sphere;

  double U(const SoftSphere<D, n>& s) const {
    double r2 = Sphere<D>::regularizedDistanceSquared(s);
    if (r2 < pow(1.25, 2)) {
      return pow(1 / r2, n / 2) + c0 + c2 * r2 + c4 * pow(r2, 2);
    } else {
      return 0;
    }
  }
};

template <int D> class Euclidean {
public:
  Vector<D> t;
  Matrix<D> r;

  Euclidean() {
    t.setZero();
    r.setIdentity();
  }

  Euclidean(const Vector<D>& t0, const Matrix<D>& r0) : t(t0), r(r0) {}

  Euclidean inverse() const { return Euclidean(-r.transpose() * t, r.transpose()); }

  Vector<D> act(const Vector<D>& x) const { return r * x; }
  Position<D> act(const Position<D>& x) const { return (Position<D>)(r * x) + t; }
};

template <int D, Particle<D> T> class Model {
private:
  NeighborLists<D, T> nl;

public:
  std::vector<T> particles;

  Model(unsigned N, double l) : nl(l) {
    particles.reserve(N);
  };

  void insert(const T& p) { 
    particles.push_back(p);
    nl.insert(particles.back());
  }

  void move(T& p, const Position<D>& x) {
    nl.move(p, x);
    p.x = x;
  }

  std::set<T*> neighborsOf(const Position<D>& x) const {
    return nl.neighborsOf(x);
  }

  double energyChange(const T& p, const T& pNew) const {
    double ΔE = 0;
    for (const T* neighbor : neighborsOf(p.x)) {
      if (neighbor != &p) {
        ΔE -= neighbor->U(p);
      }
    }
    for (const T* neighbor : neighborsOf(pNew.x)) {
      if (neighbor != &p) {
        ΔE += neighbor->U(pNew);
      }
    }

    return ΔE;
  }

  double energyChangeSwap(const T& p1, const T& p2) const {
    double ΔE = 0;
    T p1New = p1;
    T p2New = p2;
    p1New.r = p2.r;
    p2New.r = p1.r;
    for (const T* neighbor : neighborsOf(p1.x)) {
      if (neighbor != &p1 && neighbor != &p2) {
        ΔE -= neighbor->U(p1);
        ΔE += neighbor->U(p1New);
      }
    }
    for (const T* neighbor : neighborsOf(p2.x)) {
      if (neighbor != &p1 && neighbor != &p2) {
        ΔE -= neighbor->U(p2);
        ΔE += neighbor->U(p2New);
      }
    }

    return ΔE;
  }

  bool metropolis(double β, T& p, const Vector<D>& Δx, Rng& r) {
    T pNew = p;
    pNew.x += Δx;

    double ΔE = energyChange(p, pNew);

    if (exp(-β * ΔE) > r.uniform(0.0, 1.0)) {
      move(p, pNew.x);
      return true;
    } else {
      return false;
    }
  }

  bool swap(double β, T& p1, T& p2, Rng& r) {
    double ΔE = energyChangeSwap(p1, p2);
    if (exp(-β * ΔE) > r.uniform(0.0, 1.0)) {
      std::swap(p1.r, p2.r);
      return true;
    } else {
      return false;
    }
  }

  unsigned clusterFlip(double β, const Euclidean<D>& R, T& p0, Rng& r) {
    unsigned n = 0;
    std::queue<T*> q;
    q.push(&p0);

    while (!q.empty()) {
      T* p = q.front();
      q.pop();

      if (!p->m) {
        p->m = true;
        T pNew = *p;
        pNew.x = R.act(p->x);

        for (T* neighbor : neighborsOf(pNew.x)) {
          if (!neighbor->m) {
            double ΔE = neighbor->U(pNew) - neighbor->U(*p);
            if (1 - exp(-β * ΔE) > r.uniform(0.0, 1.0)) {
              q.push(neighbor);
            }
          }
        }

        move(*p, pNew.x);
        n++;
      }
    }

    for (T& p : particles) {
      p.m = false;
    }

    return n;
  }
};