#include <list>
#include <queue>

#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/src/Core/CwiseNullaryOp.h>

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

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

using Rng = randutils::random_generator<pcg32>;

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 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 < 0) ? (a + (1 - a / (signed)b) * b) : a) % b; }

template <int D, typename Derived> Vector<D> mod(const Eigen::MatrixBase<Derived>& v, unsigned b) {
  Vector<D> u;
  for (unsigned i = 0; i < D; i++) {
    u(i) = mod(v(i), b);
  }
  return u;
}

template <int D> void one_sequences(std::list<std::array<double, D>>& sequences, unsigned level) {
  if (level > 0) {
    unsigned new_level = level - 1;
    unsigned old_length = sequences.size();
    for (std::array<double, D>& sequence : sequences) {
      std::array<double, D> new_sequence = sequence;
      new_sequence[new_level] = -1;
      sequences.push_front(new_sequence);
    }
    one_sequences<D>(sequences, new_level);
  }
}

template <unsigned D> std::vector<Matrix<D>> generateTorusMatrices() {
  std::vector<Matrix<D>> mats;

  std::array<double, D> ini_sequence;
  ini_sequence.fill(1);
  std::list<std::array<double, D>> sequences;
  sequences.push_back(ini_sequence);

  one_sequences<D>(sequences, D);

  sequences.pop_back(); // don't want the identity matrix!

  for (std::array<double, D> sequence : sequences) {
    Matrix<D> m;
    for (unsigned i = 0; i < D; i++) {
      for (unsigned j = 0; j < D; j++) {
        if (i == j) {
          m(i, j) = sequence[i];
        } else {
          m(i, j) = 0;
        }
      }
    }

    mats.push_back(m);
  }

  for (unsigned i = 0; i < D; i++) {
    for (unsigned j = 0; j < D; j++) {
      if (i != j) {
        Matrix<D> m;
        for (unsigned k = 0; k < D; k++) {
          for (unsigned l = 0; l < D; l++) {
            if ((k == i && l == j) || (k == j && l == i)) {
              if (i < j) {
                m(k, l) = 1;
              } else {
                m(k, l) = -1;
              }
            } else if (k == l && (k != i && k != j)) {
              m(k, l) = 1;
            } else {
              m(k, l) = 0;
            }
          }
        }
        mats.push_back(m);
      }
    }
  }

  return mats;
}

template <unsigned D> class Transformation {
public:
  unsigned L;
  Matrix<D> m;
  Vector<D> v;

  Transformation(unsigned L) : L(L) {
    m.setIdentity();
    v.setZero();
  }

  Transformation(unsigned L, const Matrix<D>& m, const Vector<D>& v) : L(L), m(m), v(v) {}

  Transformation(unsigned L, const std::vector<Matrix<D>>& ms, Rng& r) : m(r.pick(ms)), L(L) {
    for (unsigned i = 0; i < D; i++) {
      v[i] = r.uniform((unsigned)0, L - 1);
    }

    v = v - m * v;
  }

  Vector<D> apply(const Vector<D>& x) const { return mod<D>(v + m * x, L); }

  Transformation<D> apply(const Transformation<D>& t) const {
    Transformation<D> tNew(L);

    tNew.m = m * t.m;
    tNew.v = apply(t.v);

    return tNew;
  }

  Transformation<D> inverse() const {
    return Transformation<D>(L, m.transpose(), -m.transpose() * v);
  }
};

template <typename T> concept State = requires(T v, const T& v2) {
  { v.empty() } -> std::same_as<bool>;
  {v.remove()};
  { v.operator==(v2) } -> std::same_as<bool>;
};

template <unsigned D, State S> class HalfEdge;

template <unsigned D, State S> class Vertex {
public:
  Vector<D> position;
  Vector<D> initialPosition;
  S state;
  std::vector<HalfEdge<D, S>> adjacentEdges;
  bool marked;

  bool empty() const { return state.empty(); }
};

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

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

template <unsigned D, State S> class System {
public:
  const unsigned L;
  unsigned N;
  std::vector<Vertex<D, S>> vertices;

  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;
  }

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

    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, S>(vertices[j], Δx));
        vertices[j].adjacentEdges.push_back(HalfEdge<D, S>(vertices[i], -Δx));
      }
    }
  }

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

  void setInitialPosition() {
    for (Vertex<D, S>& v : vertices) {
      v.initialPosition = v.position;
    }
  }

  double selfIntermediateScattering() const {
    double F = 0;

    Vector<D> k1 = {M_PI, 0};
    Vector<D> k2 = {0, M_PI};
    for (const Vertex<D, S>& v : vertices) {
      if (!v.empty()) {
        F += cos(k1.dot(v.position - v.initialPosition));
        F += cos(k2.dot(v.position - v.initialPosition));
      }
    }

    return F / (2 * this->N);
  }
};

template <unsigned D, State S> class FiniteEnergySystem : public System<D, S> {
public:
  using System<D, S>::System;

  virtual double pairEnergy(const S&, const S&) const { return 0; }

  double siteEnergy(const Vertex<D, S>& v) const {
    double E = 0;

    for (const HalfEdge<D, S>& e : v.adjacentEdges) {
      E += pairEnergy(v.state, e.neighbor.state);
    }

    return E;
  }

  double energy() const {
    double E = 0;

    for (const Vertex<D, S>& v : this->vertices) {
      E += siteEnergy(v);
    }

    return E / 2;
  }

  bool trySwap(Vertex<D, S>& v1, Vertex<D, S>& v2, double T, Rng& r) {
    double E0 = siteEnergy(v1) + siteEnergy(v2);
    std::swap(v1.state, v2.state);
    double E1 = siteEnergy(v1) + siteEnergy(v2);
    double ΔE = E1 - E0;

    if (exp(-ΔE / T) > r.uniform(0.0, 1.0)) {
      /* Accept the swap. */
      std::swap(v1.initialPosition, v2.initialPosition);
      return true;
    } else {
      /* Revert the swap. */
      std::swap(v1.state, v2.state);
      return false;
    }
  }

  bool tryRandomMove(double T, Rng& r) {
    Vertex<D, S>& v1 = r.pick(this->vertices);

    if (v1.empty()) {
      return false;
    }

    Vertex<D, S>& v2 = (r.pick(v1.adjacentEdges)).neighbor;

    if (!v2.empty()) {
      return false;
    }

    return trySwap(v1, v2, T, r);
  }

  bool tryRandomSwap(double T, Rng& r) {
    Vertex<D, S>& v1 = r.pick(this->vertices);
    Vertex<D, S>& v2 = r.pick(this->vertices);

    if (v1.state != v2.state) {
      return trySwap(v1, v2, T, r);
    } else {
      return false;
    }
  }

  void sweepLocal(Rng& r) {
    for (unsigned i = 0; i < iPow(this->L, D); i++) {
      tryRandomMove(r);
    }
  }

  void sweepSwap(Rng& r) {
    for (unsigned i = 0; i < iPow(this->L, D); i++) {
      tryRandomSwap(r);
    }
  }

  unsigned flipCluster(const Transformation<D>& R, Vertex<D, S>& v0, double T, Rng& r,
                       bool dry = false) {
    std::queue<std::array<std::reference_wrapper<Vertex<D, S>>, 2>> q;
    Vector<D> x0New = R.apply(v0.position);
    Vertex<D, S>& v0New = this->vertices[this->vectorToIndex(x0New)];
    q.push({v0, v0New});

    unsigned n = 0;

    while (!q.empty()) {
      auto [vR, vNewR] = q.front();
      q.pop();

      Vertex<D, S>& v = vR;
      Vertex<D, S>& vNew = vNewR;

      if (!v.marked && !vNew.marked) {
        v.marked = true;
        vNew.marked = true;

        for (HalfEdge<D, S>& e : v.adjacentEdges) {
          Vertex<D, S>& vn = e.neighbor;

          Vector<D> xnNew = R.apply(vn.position);
          Vertex<D, S>& vnNew = this->vertices[this->vectorToIndex(xnNew)];

          if (!vn.marked && !vnNew.marked) {
            double E0 = pairEnergy(v.state, vn.state) + pairEnergy(vNew.state, vnNew.state);
            double E1 = pairEnergy(vNew.state, vn.state) + pairEnergy(v.state, vnNew.state);
            double ΔE = E1 - E0;

            if (exp(-ΔE / T) < r.uniform(0.0, 1.0)) {
              q.push({vn, vnNew});
            }
          }
        }

        if (!dry) {
          std::swap(v.state, vNew.state);
          std::swap(v.initialPosition, vNew.initialPosition);
        }

        n += 1;
      }
    }

    return n;
  }

  unsigned wolff(const Transformation<D>& R, double T, Rng& r) {
    unsigned n = flipCluster(R, r.pick(this->vertices), T, r);
    for (Vertex<D, S>& v : this->vertices) {
      v.marked = false;
    }
    return n;
  }

  void swendsenWang(const Transformation<D>& R, double T, Rng& r) {
    for (Vertex<D, S>& v : this->vertices) {
      if (!v.marked) {
        bool dry = 0.5 < r.uniform(0.0, 1.0);
        flipCluster(R, v, T, r, dry);
      }
    }

    for (Vertex<D, S>& v : this->vertices) {
      v.marked = false;
    }
  }
};