#include <iostream>
#include <queue>

#include "glass.hpp"

class DistinguishableState {
public:
  unsigned type;

  DistinguishableState() { type = 0; }

  DistinguishableState(unsigned t) { type = t; }

  bool empty() const { return type == 0; }
  void remove() { type = 0; }
};

template <unsigned D> class System {
public:
  const unsigned L;
  unsigned N;
  std::vector<Vertex<D, DistinguishableState>> vertices;
  std::vector<double> interaction;

  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, Rng& r) : L(L), N(N), vertices(iPow(L, D)), interaction(N * N) {
    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, DistinguishableState>(vertices[j], Δx));
        vertices[j].adjacentEdges.push_back(HalfEdge<D, DistinguishableState>(vertices[i], -Δx));
      }
    }

    for (double& V : interaction) {
      V = r.uniform(-0.5, 0.5);
    }

    for (unsigned i = 0; i < N; i++) {
      vertices[i].state.type = i + 1;
    }
  }

  double pairEnergy(const DistinguishableState& s1, const DistinguishableState& s2) const {
    if (s1.empty() || s2.empty()) {
      return 0;
    } else {
      if (s1.type < s2.type) {
        return interaction[(s1.type - 1) * N + (s2.type - 1)];
      } else {
        return interaction[(s2.type - 1) * N + (s1.type - 1)];
      }
    }
  }

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

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

    return E;
  }

  double energy() const {
    double E = 0;

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

    return E / 2;
  }

  bool trySwap(Vertex<D, DistinguishableState>& v1, Vertex<D, DistinguishableState>& 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, DistinguishableState>& v1 = r.pick(vertices);

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

    Vertex<D, DistinguishableState>& 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, DistinguishableState>& v1 = r.pick(vertices);
    Vertex<D, DistinguishableState>& v2 = r.pick(vertices);

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

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

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

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

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

    unsigned n = 0;

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

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

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

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

          Vector<D> xnNew = R.apply(vn.position);
          Vertex<D, DistinguishableState>& vnNew = vertices[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(vertices), T, r);
    for (Vertex<D, DistinguishableState>& v : vertices) {
      v.marked = false;
    }
    return n;
  }

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

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

  int overlap(const System<D>& s) const {
    int o = 0;

    for (unsigned i = 0; i < vertices.size(); i++) {
      //      DistinguishableState<D> s2 =
      //      orientation.apply(s.vertices[vectorToIndex(orientation.inverse().apply(indexToVector(i)))].state);
      //    o += vertices[i].state.dot(s2);
    }

    return o;
  }

  void setInitialPosition() {
    for (Vertex<D, DistinguishableState>& 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, DistinguishableState>& v : vertices) {
      if (!v.empty()) {
        F += cos(k1.dot(v.position - v.initialPosition));
        F += cos(k2.dot(v.position - v.initialPosition));
      }
    }

    return F / (2 * N);
  }
};

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 = 2;
  unsigned L = 40;
  unsigned N = 1560;
  double Tmin = 0;
  double Tmax = 0.4;
  double δT = 0.02;

  Rng r;

  System<D> s(L, N, r);

  std::vector<Matrix<D>> ms = generateTorusMatrices<D>();

  for (unsigned j = 0; j < 10 * s.vertices.size(); j++) {
    // unsigned nC = s.wolff(Transformation<D>(L, m, v.position - m * v.position), T, r);
    s.tryRandomSwap(1000, r);
  }

  for (unsigned i = 0; i < 1e5; i++) {
    for (unsigned j = 0; j < s.vertices.size(); j++) {
      s.tryRandomMove(Tmax, r);
    }
  }

  unsigned n = 1;
  for (double T = Tmax; T > Tmin; T -= δT) {

    /*
    for (unsigned i = 0; i < 1e5; i++) {
      Matrix<D> m = r.pick(ms);
      Vertex<D, DistinguishableState>& v = r.pick(s.vertices);
      s.wolff(Transformation<D>(L, m, v.position - m * v.position), T, r);
    }
    */
    /*
     */
    s.setInitialPosition();
    std::cout << T << " ";
    for (unsigned i = 0; i < 1e4; i++) {
      Matrix<D> m = r.pick(ms);
      Vertex<D, DistinguishableState>& v = r.pick(s.vertices);
      s.wolff(Transformation<D>(L, m, v.position - m * v.position), T, r);
      std::cout << s.selfIntermediateScattering() << " ";
    }
    /*
    for (unsigned i = 0; i < 1e5; i++) {
      for (unsigned j = 0; j < s.vertices.size(); j++) {
        s.tryRandomMove(T, r);
      }
      std::cout << s.selfIntermediateScattering() << " ";
    }
    */
    std::cout << std::endl;
    std::cerr << T << " " << s.energy() / N << std::endl;
    //    s.sweepLocal(r);
    //    s.sweepSwap(r);
    //    s.swendsenWang(Transformation<D>(L, ms, r), r);
  }

  return 0;
}