#pragma once

#include <array>
#include <functional>
#include <list>
#include <queue>
#include <vector>

#include "randutils/randutils.hpp"

#include "euclidean.hpp"
#include "octree.hpp"
#include "quantity.hpp"
#include "spin.hpp"

template <class U, int D, class R, class S> class Model;

template <class U, int D, class R, class S> class measurement {
public:
  virtual void pre_cluster(const Model<U, D, R, S>&, unsigned, const R&){};
  virtual void plain_bond_visited(const Model<U, D, R, S>&, const Spin<U, D, S>*,
                                  const Spin<U, D, S>*, const Spin<U, D, S>&, double){};
  virtual void plain_site_transformed(const Model<U, D, R, S>&, const Spin<U, D, S>*,
                                      const Spin<U, D, S>&){};

  virtual void ghost_bond_visited(const Model<U, D, R, S>&, const Spin<U, D, S>&,
                                  const Spin<U, D, S>&, double){};
  virtual void ghost_site_transformed(const Model<U, D, R, S>&, const R&){};

  virtual void post_cluster(const Model<U, D, R, S>&){};
};

template <class U, int D, class R, class S> class Model {
public:
  U L;
  R s0;
  std::vector<Spin<U, D, S>> s;
  Octree<U, D, S> dict;
  std::function<double(const Spin<U, D, S>&, const Spin<U, D, S>&)> Z;
  std::function<double(const Spin<U, D, S>&)> B;
  randutils::mt19937_rng rng;

  Model(U L, std::function<double(const Spin<U, D, S>&, const Spin<U, D, S>&)> Z,
        std::function<double(const Spin<U, D, S>&)> B)
      : L(L), s0(L), Z(Z), B(B), rng(), dict(L, std::floor(L)) {}

  void step(double T, unsigned ind, R& r, measurement<U, D, R, S>& A) {
    std::queue<Spin<U, D, S>*> queue;
    queue.push(&(s[ind]));

    std::set<Spin<U, D, S>*> visited;

    while (!queue.empty()) {
      Spin<U, D, S>* si = queue.front();
      queue.pop();

      if (!visited.contains(si)) {
        visited.insert(si);

        if (si == NULL) {
          R s0_new = r.act(s0);
          for (Spin<U, D, S>& ss : s) {
            Spin<U, D, S> s0s_old = s0.inverse().act(ss);
            Spin<U, D, S> s0s_new = s0_new.inverse().act(ss);
            double p = 1.0 - exp(-(B(s0s_new) - B(s0s_old)) / T);
            A.ghost_bond_visited(*this, s0s_old, s0s_new, p);
            if (rng.uniform(0.0, 1.0) < p) {
              queue.push(&ss);
            }
          }
          A.ghost_site_transformed(*this, s0_new);
          s0 = s0_new;
        } else {
          Spin<U, D, S> si_new = r.act(*si);
          std::set<Spin<U, D, S>*> all_neighbors;
          std::set<Spin<U, D, S>*> current_neighbors = dict.neighbors(si->x);
          std::set<Spin<U, D, S>*> new_neighbors = dict.neighbors(si_new.x);

          all_neighbors.insert(current_neighbors.begin(), current_neighbors.end());
          all_neighbors.insert(new_neighbors.begin(), new_neighbors.end());
          all_neighbors.insert(NULL);
          for (Spin<U, D, S>* sj : all_neighbors) {
            if (sj != si) {
              double p;
              if (sj == NULL) {
                Spin<U, D, S> s0s_old = s0.inverse().act(*si);
                Spin<U, D, S> s0s_new = s0.inverse().act(si_new);
                p = 1.0 - exp(-(B(s0s_new) - B(s0s_old)) / T);
                A.ghost_bond_visited(*this, s0s_old, s0s_new, p);
              } else {
                p = 1.0 - exp(-(Z(*si, *sj) - Z(si_new, *sj)) / T);
                A.plain_bond_visited(*this, si, sj, si_new, p);
              }
              if (rng.uniform(0.0, 1.0) < p) {
                queue.push(sj);
              }
            }
          }
          A.plain_site_transformed(*this, si, si_new);
          dict.remove(si);
          *si = si_new;
          dict.insert(si);
        }
      }
    }
  }

  void resize(double T, double P) {}

  void wolff(double T, std::function<R(const Model<U, D, R, S>&, randutils::mt19937_rng&)> g,
             measurement<U, D, R, S>& A, unsigned N) {
    for (unsigned i = 0; i < N; i++) {
      R r = g(*this, rng);
      unsigned ind = rng.uniform((unsigned)0, (unsigned)(s.size() - 1));

      A.pre_cluster(*this, ind, r);

      this->step(T, ind, r, A);

      A.post_cluster(*this);
    }
  }
};