#pragma once

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

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

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

using Rng = randutils::random_generator<pcg32>;

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 Transformation {
  public:
    virtual std::set<Spin<U, D, S>*> current() const {return {NULL};}
    virtual std::set<Spin<U, D, S>*> toConsider() const {return {};}
    virtual double ΔE(const Spin<U, D, S>*) const {return 0;}
    virtual Transformation* createNew(Spin<U, D, S>*) const {return new Transformation();}
    virtual void apply();
};

template <class U, int D, class R, class S> using Gen = std::function<Transformation<U, D, R, S>*(Model<U, D, R, S>&, Rng&)>;

template <class U, int D, class R, class S> class SpinFlip : public Transformation<U, D, R, S> {
  private:
    Model<U, D, R, S>& M;
    Spin<U, D, S>& sOld;
    const R r;
    Spin<U, D, S>  sNew;
  public:
    SpinFlip(Model<U, D, R, S>& M, const R& r, Spin<U, D, S>& s) : M(M), r(r), sOld(s) {
      sNew = r.act(s);
    }

    std::set<Spin<U, D, S>*> current() const override {
      return {&sOld};
    }

    std::set<Spin<U, D, S>*> toConsider() const override {
      std::set<Spin<U, D, S>*> neighbors;
      std::set<Spin<U, D, S>*> current_neighbors = M.dict.neighbors(sOld.x);
      std::set<Spin<U, D, S>*> new_neighbors = M.dict.neighbors(sNew.x);

      neighbors.insert(current_neighbors.begin(), current_neighbors.end());
      neighbors.insert(new_neighbors.begin(), new_neighbors.end());
      neighbors.insert(NULL);

      return neighbors;
    }

    double ΔE(const Spin<U, D, S>* s) const override {
      if (s == NULL) {
        Spin<U, D, S> s0s_old = M.s0.inverse().act(sOld);
        Spin<U, D, S> s0s_new = M.s0.inverse().act(sNew);
        return M.B(s0s_new) - M.B(s0s_old);
      } else {
        return M.Z(sOld, *s) - M.Z(sNew, *s);
      }
    }

    Transformation<U, D, R, S>* createNew(Spin<U, D, S>* s) const override;

    void apply() override {
      M.dict.remove(&sOld);
      sOld = sNew;
      M.dict.insert(&sOld);
    }
};

template <class U, int D, class R, class S> class PairFlip : public Transformation<U, D, R, S> {
  private:
    Model<U, D, R, S>& M;
    Spin<U, D, S>& s1Old;
    Spin<U, D, S>& s2Old;
    const R r;
    Spin<U, D, S>  s1New;
    Spin<U, D, S>  s2New;
  public:
    PairFlip(Model<U, D, R, S>& M, const R& r, Spin<U, D, S>& s1, Spin<U, D, S>& s2) : M(M), r(r), s1Old(s1), s2Old(s2) {
      s1New = r.act(s1);
      s2New = r.act(s2);
    }

    std::set<Spin<U, D, S>*> current() const override {
      return {&s1Old, &s2Old};
    }

    std::set<Spin<U, D, S>*> toConsider() const override {
      std::set<Spin<U, D, S>*> neighbors;
      std::set<Spin<U, D, S>*> current_neighbors_1 = M.dict.neighbors(s1Old.x);
      std::set<Spin<U, D, S>*> current_neighbors_2 = M.dict.neighbors(s2Old.x);
      std::set<Spin<U, D, S>*> new_neighbors_1 = M.dict.neighbors(s1New.x);
      std::set<Spin<U, D, S>*> new_neighbors_2 = M.dict.neighbors(s2New.x);

      neighbors.insert(current_neighbors_1.begin(), current_neighbors_1.end());
      neighbors.insert(current_neighbors_2.begin(), current_neighbors_2.end());
      neighbors.insert(new_neighbors_1.begin(), new_neighbors_1.end());
      neighbors.insert(new_neighbors_2.begin(), new_neighbors_2.end());
      neighbors.insert(NULL);

      return neighbors;
    }

    double ΔE(const Spin<U, D, S>* s) const override {
      if (s == NULL) {
        Spin<U, D, S> s0s1_old = M.s0.inverse().act(s1Old);
        Spin<U, D, S> s0s1_new = M.s0.inverse().act(s1New);
        Spin<U, D, S> s0s2_old = M.s0.inverse().act(s2Old);
        Spin<U, D, S> s0s2_new = M.s0.inverse().act(s2New);
        return M.B(s0s1_new) + M.B(s0s2_new) - M.B(s0s1_old) - M.B(s0s2_old);
      } else {
        return M.Z(s1Old, *s) + M.Z(s2Old, *s) - M.Z(s1New, *s) - M.Z(s2New, *s);
      }
    }

    Transformation<U, D, R, S>* createNew(Spin<U, D, S>* s) const override;

    void apply() override {
      M.dict.remove(&s1Old);
      M.dict.remove(&s2Old);
      s1Old = s1New;
      s2Old = s2New;
      M.dict.insert(&s1Old);
      M.dict.insert(&s2Old);
    }
};

template <class U, int D, class R, class S> class FieldFlip : public Transformation<U, D, R, S> {
  private:
    Model<U, D, R, S>& M;
    const R r;
    R s0New;
  public:
    FieldFlip(Model<U, D, R, S>& M, const R& r) : M(M), r(r), s0New(r.act(M.s0)) {}

    std::set<Spin<U, D, S>*> toConsider() const override {
      std::set<Spin<U, D, S>*> neighbors;

      for (Spin<U, D, S>& s : M.s) {
        neighbors.insert(&s);
      }

      return neighbors;
    }

    double ΔE(const Spin<U, D, S>* s) const override {
      Spin<U, D, S> s0s_old = M.s0.inverse().act(*s);
      Spin<U, D, S> s0s_new = s0New.inverse().act(*s);
      return M.B(s0s_new) - M.B(s0s_old);
    }

    Transformation<U, D, R, S>* createNew(Spin<U, D, S>* s) const override {
      return new SpinFlip<U, D, R, S>(M, r, *s);
    }

    void apply() override {
      M.s0 = s0New;
    }
};

template<class U, int D, class R, class S> Transformation<U, D, R, S>* SpinFlip<U, D, R, S>::createNew(Spin<U, D, S>* s) const {
  if (s == NULL) {
    return new FieldFlip<U, D, R, S>(M, r);
  } else {
    return new SpinFlip(M, r, *s);
  }
}

template<class U, int D, class R, class S> Transformation<U, D, R, S>* PairFlip<U, D, R, S>::createNew(Spin<U, D, S>* s) const {
  if (s == NULL) {
    return new FieldFlip<U, D, R, S>(M, r);
  } else {
    return new SpinFlip<U, D, R, S>(M, 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;
  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, Transformation<U, D, R, S>* t0, measurement<U, D, R, S>& A) {
    std::queue<Transformation<U, D, R, S>*> queue;
    queue.push(t0);

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

    while (!queue.empty()) {
      Transformation<U, D, R, S>* t = queue.front();
      queue.pop();

      std::set<Spin<U, D, S>*> c = t->current();
      if (!std::any_of(c.begin(), c.end(), [&visited](Spin<U, D, S>* s) { return visited.contains(s); })) {
        visited.insert(c.begin(), c.end());

        for (Spin<U, D, S>* s : t->toConsider()) {
          double ΔE = t->ΔE(s);
          if (ΔE > 0 && 1.0 - exp(-ΔE / T) > rng.uniform<double>(0, 1)) {
            queue.push(t->createNew(s));
          }
        }

        t->apply();
      }
      delete t;
    }
  }

  void resize(double T, double P) {}

  void wolff(double T, std::vector<Gen<U, D, R, S>> gs,
             measurement<U, D, R, S>& A, unsigned N) {
    for (unsigned i = 0; i < N; i++) {
      Gen<U, D, R, S>& g = rng.pick(gs);
      Transformation<U, D, R, S> *t = g(*this, rng);

      this->step(T, t, A);

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