#include <iostream>
#include <limits>
#include <vector>
#include <list>
#include <set>
#include <eigen3/Eigen/Dense>
#include <unordered_map>

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

using Rng = randutils::random_generator<pcg32>;

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

// Integer power function
int iPow(int x, unsigned p) {
  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 mod(signed a, unsigned b) {
  return ((a %= b) < 0) ? a + b : a;
}

template<unsigned D> class State {
  public:
    Vector<D> σ;

    State(unsigned a, signed b) : σ(Vector<D>::Zero()) {
      σ[a] = b;
    }

    State(Rng& r) : State(r.uniform((unsigned)0, D - 1), r.pick({-1, 1})) {}

    State<D> apply(const Matrix<D>& m) const {
      return {m * σ};
    }
};

template <unsigned D> class Vertex;

template<unsigned D> class Particle {
  public:
    Vertex<D>& position;
    State<D> state;
    bool marked;

    Particle(Vertex<D>& p, const State<D>& s) : position(p), state(s) {}
};

template<unsigned D> class HalfEdge {
  public:
    Vertex<D>& neighbor;
    Vector<D> Δx;

    HalfEdge(Vertex<D>& n, const Vector<D>& d) : neighbor(n), Δx(d) {}
};

template<unsigned D> class Vertex {
  public:
    Vector<D> position;
    std::vector<HalfEdge<D>> adjacentEdges;
    typename std::unordered_map<unsigned, Particle<D>>::iterator occupant;
};

template<unsigned D> class System {
  public:
    const unsigned L;
    std::vector<Vertex<D>> vertices;
    std::unordered_map<unsigned, Particle<D>> particles;

    unsigned vectorToIndex(const Vector<D>& x) const {
      unsigned i = 0;
      for (unsigned d = 0; d < D; d++) {
        i += x[d] * iPow(L, d);
      }
      return i;
    }

    Vector<D> indexToVector(unsigned i) const {
      Vector<D> x;
      for (unsigned d = 0; d < D; d++) {
        x[d] = (i / iPow(L, d)) % L;
      }
      return x;
    }

    bool isEmpty(const Vertex<D>& v) const {
      return v.occupant == particles.end();
    }

    unsigned N() const {
      return particles.size();
    }

    System(unsigned L) : L(L), vertices(iPow(L, D)) {
      for (unsigned i = 0; i < iPow(L, D); i++) {
        vertices[i].position = indexToVector(i);
        vertices[i].adjacentEdges.reserve(2 * D);
        vertices[i].occupant = particles.end();
      }

      for (unsigned d = 0; d < D; d++) {
        Vector<D> Δx = Vector<D>::Zero();
        Δx[d] = 1;
        for (signed i = 0; i < iPow(L, D); i++) {
          unsigned j = iPow(L, d + 1) * (i / iPow(L, d + 1)) + mod(i + iPow(L, d), pow(L, d + 1));
          vertices[i].adjacentEdges.push_back(HalfEdge<D>(vertices[j], Δx));
          vertices[j].adjacentEdges.push_back(HalfEdge<D>(vertices[i], -Δx));
        }
      }
    }

    std::list<typename std::unordered_map<unsigned, Particle<D>>::iterator> overlaps(const Particle<D>& p, bool excludeSelf = false) const {
      std::list<typename std::unordered_map<unsigned, Particle<D>>::iterator> o;

      if (!excludeSelf && !isEmpty(p.position)) {
        o.push_back(p.position.occupant);
      }

      for (const HalfEdge<D>& e : p.position.adjacentEdges) {
        if (!isEmpty(e.neighbor)) {
          if (p.state.σ.dot(e.Δx) == 1 || e.neighbor.occupant->second.state.σ.dot(e.Δx) == -1) {
            o.push_back(e.neighbor.occupant);
          } 
        }
      }

      return o;
    }

    bool insert(const Particle<D>& p, Rng& r) {
      if (overlaps(p).empty()) {
        typename std::unordered_map<unsigned, Particle<D>>::iterator it;
        std::tie(it, std::ignore) = particles.insert({r.uniform((unsigned)0, (unsigned)pow(2, 28)), p});
        p.position.occupant = std::prev(particles.end());
        return true;
      } else {
        return false;
      }
    }

    void remove(typename std::unordered_map<unsigned, Particle<D>>::iterator ip) {
      ip->second.position.occupant = particles.end();
      particles.erase(ip);
    }

    /*
    bool randomMove(Rng& r) {
      Particle<D>& p = *(r.pick(particles));
      Particle<D> pN = p;

      if (0.2 < r.uniform(0.0, 1.0)) {
        pN.position = r.pick(p.position.adjacentEdges).neighbor;
      }
      pN.state = randomState<D>(r);

      if (overlaps(pN, true).empty()) {
        remove(&p);
        insert(pN);
        return true;
      } else {
        return false;
      }
    }
    */

    bool swap(Particle<D>& p1, Particle<D>& p2) {
      std::swap(p1.state, p2.state);
      if (overlaps(p1, true).empty() && overlaps(p2, true).empty()) {
        return true;
      } else {
        std::swap(p1.state, p2.state);
        return false;
      }
    }

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

    void setGroundState(Rng& r) {
      for (typename std::unordered_map<unsigned, Particle<D>>::iterator ip = particles.begin(); ip != particles.end(); ip++) {
        remove(ip);
      }

      for (Vertex<D>& v : vertices) {
        unsigned a = 0;
        for (unsigned d = 0; d < D; d++) {
          a += (d + 1) * v.position(d);
        }
        a %= 2 * D + 1;

        if (a > 0) {
          if (a <= D) {
            State<D> s(a - 1, -1);
            insert(Particle<D>(v, s), r);
          } else if (D < a) {
            State<D> s(2 * D - a, 1);
            insert(Particle<D>(v, s), r);
          }
        }
      }
    }

    bool compatible() {
      for (Particle<D> p : particles) {
        if (!overlaps(p, true).empty()) {
          return false;
        }
      }

      return true;
    }

    double density() const {
      return N() / pow(L, D);
    }

    unsigned maxOccupation() const {
//      return (2 * D * iPow(L, D)) / (2 * D + 1);
      return iPow(L, D);
    }

    void sweepGrandCanonical(double z, Rng& r) {
      for (unsigned i = 0; i < iPow(L, D); i++) {
        if (0.5 < r.uniform(0.0, 1.0)) {
          double pIns = maxOccupation() * z / (N() + 1);

          if (pIns > r.uniform(0.0, 1.0)) {
            while (true) {
              Vertex<D>& v = r.pick(vertices);
              if (isEmpty(v)) {
                insert(Particle<D>(v, State<D>(r)), r);
                break;
              }
            }
          }
        } else {

          double pDel = N() / (z * maxOccupation());

          if (pDel > r.uniform(0.0, 1.0)) {
            remove(r.choose(particles));
          }
        }

//        randomMove(r);
        randomSwap(r);
      }
    }

    void sweepLocal(double z, Rng& r) {
      for (unsigned i = 0; i < iPow(L, D); i++) {
//        randomMove(r);
      }
    }

    void sweepSwap(double z, Rng& r) {
      for (unsigned i = 0; i < iPow(L, D); i++) {
        randomSwap(r);
      }
    }

    void flipCluster(const Matrix<D>& R, Vertex<D>& v0, Rng& r, bool dry = false) {
      std::list<std::reference_wrapper<Vertex<D>>> q = {v0};

      while (!q.empty()) {
        Vertex<D>& v = q.front();
        q.pop_front();

        Vector<D> xNew = R * v.position;
        State<D> sNew = v.state.apply(R);

        for (Vertex<D>& vn : overlaps(vertices[vectorToIndex(xNew)], sNew)) {
          if (!vn.marked) {
            vn.marked = true;
            q.push_back(vn);
          }
        }
      }
    }
};

template<unsigned D> Vector<D> randomVector(unsigned L, Rng& r) {
  Vector<D> x;
  for (unsigned i = 0; i < D; i++) {
    x[i] = r.uniform((unsigned)0, L - 1);
  }
  return x;
}

int main() {
  const unsigned D = 3;
  unsigned L = 28;
  unsigned Nmin = 2e2;
  unsigned Nmax = 2e5;
  double Tmin = 0.04;
  double Tmax = 0.2;
  double δT = 0.02;

  System<D> s(L);

  Rng r;

  /*
  for (unsigned N = Nmin; N <= Nmax; N *= 10) {
    s.setGroundState();
    if (!s.compatible()) {
      return 1;
    }
    for (double T = Tmin; T < Tmax; T *= 1 + δT) {
      double z = exp(1 / T);

      for (unsigned n = 0; n < N; n++) {
        s.sweepGrandCanonical(z, r);
      }

      std::cout << L << " " << T << " " << N << " " << δT << " " << 1 / s.density() << std::endl;
    }

    for (double T = Tmax; T > Tmin; T /= 1 + δT) {
      double z = exp(1 / T);

      for (unsigned n = 0; n < N; n++) {
        s.sweepGrandCanonical(z, r);
      }

      std::cout << L << " " << T << " " << N << " " << δT << " " << 1 / s.density() << std::endl;
    }
  }
  */

  s.setGroundState(r);

  for (unsigned n = 0; n < 10; n++) {
    s.sweepGrandCanonical(exp(1/0.15), r);
  }

  return 0;
}