#pragma once

#include <assert.h>
#include <random>
#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 <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 + t; }
  Position<D> act(const Position<D>& x) const { return (Position<D>)(r * x) + t; }
  Euclidean<D> act(const Euclidean<D>& R) const { return {act(R.t), r * R.r}; }
};

template <typename T, int D> concept Particle = requires(T v, const Position<D>& x, double r) {
  { v.x } -> std::same_as<Position<D>>;
  { v.r } -> std::same_as<double>;
  { v.m } -> std::same_as<bool>;
  { v.U(x, r) } -> 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, Particle<D> T> class Move {
public:
  const Euclidean<D> g;
  std::vector<T*> ps;
  std::vector<Position<D>> xNews;

  Move(const Euclidean<D>& g, std::vector<T*> ps) : g(g), ps(ps) {
    xNews.reserve(ps.size());
    for (T* p : ps) {
      xNews.push_back(g.act((Position<D>)p->x));
    }
  }

  void insert(T* p) {
    ps.push_back(p);
    xNews.push_back(g.act(p->x));
  }

  bool anyMarked() const {
    for (const T* p : ps) {
      if (p->m) {
        return true;
      }
    }

    return false;
  }

  unsigned size() const {
    return ps.size();
  }

  double ΔE(const T& n) const {
    double e = 0;

    for (unsigned i = 0; i < ps.size(); i++) {
      if (&n == ps[i]) {
        return 0;
      } else {
        e += n.U(xNews[i], ps[i]->r) - n.U(ps[i]->x, ps[i]->r);
      }
    }

    return e;
  }
};

template <int D, Particle<D> T> class Model {
private:
  bool metropolisAccept(double β, double ΔE, Rng& r, double a = 1.0) const {
    return ΔE < 0 || exp(-β * ΔE) > r.uniform(0.0, 1.0);
  }

  bool swapAccept(const T& p1, const T& p2) const {
    double rat = p1.r / p2.r;
    return (rat < 1.2 && rat > 0.83);
  }

public:
  NeighborLists<D, T> nl;
  std::vector<T> particles;
  Euclidean<D> orientation;

  Model(unsigned N, double l, std::function<double(Rng& r)> g, Rng& r) : nl(l), particles(N) {
    for (T& p : particles) {
      for (double& xi : p.x) {
        xi = r.uniform(0.0, 1.0);
      }
      p.r = g(r);
      nl.insert(p);
    }
  };

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

  void execute(Move<D, T>& m) {
    for (unsigned i = 0; i < m.size(); i++) {
      move(*(m.ps[i]), m.xNews[i]);
      m.ps[i]->m = true;
    }
  }

  bool move(double β, T& p, const Vector<D>& Δx, Rng& r) {
    Position<D> xNew = p.x + Δx;
    double ΔE = 0;

    for (const T* neighbor : nl.neighborsOf(p.x)) {
      if (neighbor != &p) {
        ΔE -= neighbor->U(p.x, p.r);
      }
    }

    for (const T* neighbor : nl.neighborsOf(xNew)) {
      if (neighbor != &p) {
        ΔE += neighbor->U(xNew, p.r);
      }
    }

    if (metropolisAccept(β, ΔE, r)) {
      move(p, xNew);
      return true;
    } else {
      return false;
    }
  }

  bool swap(double β, T& p1, T& p2, Rng& r) {
    if (!swapAccept(p1, p2))
      return false;

    double ΔE = 0;

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

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

    if (metropolisAccept(β, ΔE, r)) {
      std::swap(p1.r, p2.r);
      return true;
    } else {
      return false;
    }
  }

  bool randomSwap(double β, Rng& r) {
    return swap(β, r.pick(particles), r.pick(particles), r);
  }

  bool randomMove(double β, double δ, Rng& r) {
    Vector<D> Δx;
    for (double& Δxi : Δx) {
      Δxi = r.variate<double, std::normal_distribution>(0, δ);
    }
    return move(β, r.pick(particles), Δx, r);
  }

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

    while (!q.empty()) {
      Move<D, T> m = q.front();
      q.pop();

      if (!m.anyMarked()) {
        std::set<T*> neighbors = {};
        for (unsigned i = 0; i < m.ps.size(); i++) {
          std::set<T*> pNeighbors = nl.neighborsOf(m.ps[i]->x);
          std::set<T*> pNeighborsNew = nl.neighborsOf(m.xNews[i]);
          neighbors.insert(pNeighbors.begin(), pNeighbors.end());
          neighbors.insert(pNeighborsNew.begin(), pNeighborsNew.end());
        }
        double ΔEmax = 0;
        T* pMax = NULL;

        for (T* neighbor : neighbors) {
          if (!neighbor->m) {
            double ΔE = m.ΔE(*neighbor);
            if (ΔE > ΔEmax) {
              ΔEmax = ΔE;
              pMax = neighbor;
            }
          }
        }

        if (pMax != NULL) {
          if (1 - exp(-β * ΔEmax) > r.uniform(0.0, 1.0)) {
            m.insert(pMax);
          }
        }
        
        for (T* neighbor : neighbors) {
          if (!neighbor->m) {
            double ΔE = m.ΔE(*neighbor);
            if (1 - exp(-β * ΔE) > r.uniform(0.0, 1.0)) {
              q.push(Move<D, T>(R, {neighbor}));
            }
          }
        }

        execute(m);
        n += m.size();;
      }
    }

    if (n > particles.size() / 2) {
      orientation = R.act(orientation);
    }

    return n;
  }

  unsigned clusterSwap(double β, T& s1, T& s2, Rng& r) {
    if (!swapAccept(s1, s2))
      return 0;

    Vector<2> t1 = s1.x;
    Vector<2> t2 = s2.x;
    Vector<2> t = (t1 + t2) / 2;

    Matrix<2> mat;

    mat(0, 0) = -1;
    mat(1, 1) = -1;
    mat(0, 1) = 0;
    mat(1, 0) = 0;

    Euclidean<2> g(t - mat * t, mat);

    Move<D, T> m(g, {&s1, &s2});

    return clusterFlip(β, g, m, r);
  }

  unsigned randomCluster(double β, Rng& r) {
    Vector<2> t = {r.uniform(0.0, 1.0), r.uniform(0.0, 1.0)};

    Matrix<2> mat;

    mat(0, 0) = -1;
    mat(1, 1) = -1;
    mat(0, 1) = 0;
    mat(1, 0) = 0;

    Euclidean<2> g(t - mat * t, mat);

    Move<D, T> m(g, {&r.pick(particles)});

    return clusterFlip(β, g, m, r);
  }

  unsigned clusterNudge(double β, T& s, Rng& r) {
    Vector<2> t = s.x + (Vector<2>){r.variate<double, std::normal_distribution>(0, s.r / 100),
        r.variate<double, std::normal_distribution>(0, s.r / 100)}
      ;

    Matrix<2> mat;

    mat(0, 0) = -1;
    mat(1, 1) = -1;
    mat(0, 1) = 0;
    mat(1, 0) = 0;

    Euclidean<2> g(t - mat * t, mat);

    Move<D, T> m(g, {&s});

    return clusterFlip(β, g, m, r);
  }

  unsigned clusterNeighborSwap(double β, T& s, Rng& r) {
    T* s2 = &s;
    while (s2 == &s) {
      s2 = r.pick(nl.neighborsOf(s.x));
    }
    return clusterSwap(β, s, *s2, r);
  }

  unsigned randomClusterNeighborSwap(double β, Rng& r) {
    return clusterNeighborSwap(β, r.pick(particles), r);
  }

  unsigned randomClusterNudge(double β, Rng& r) {
    return clusterNudge(β, r.pick(particles), r);
  }

  unsigned randomClusterSwap(double β, Rng& r) {
    return clusterSwap(β, r.pick(particles), r.pick(particles), r);
  }
};

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

  Sphere() : m(false) {};

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

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

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

std::function<double(Rng&)> gContinuousPolydisperse(double Rmax) {
  return [Rmax] (Rng& r) -> double {
    return pow(r.uniform(pow(2.219, -2), 1.0), -0.5) * Rmax / 2.219;
  };
}