#pragma once

#include <set>

#include "random.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 Transformation {
public:
  Model<U, D, R, S>& M;
  const R r;

  Transformation(Model<U, D, R, S>& M, const R& r) : r(r), M(M) {}
  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(double T, std::set<Spin<U, D, S>*>&, Spin<U, D, S>*,
                                    Rng& rng) const {
    return new Transformation(M, r);
  }
  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 FieldFlip : public Transformation<U, D, R, S> {
private:
  R s0New;

public:
  FieldFlip(Model<U, D, R, S>& M, const R& r)
      : Transformation<U, D, R, S>(M, 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 : this->M.s) {
      neighbors.insert(s);
    }

    return neighbors;
  }

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

  Transformation<U, D, R, S>* createNew(double T, std::set<Spin<U, D, S>*>& cluster,
                                             Spin<U, D, S>* s, Rng& rng) const override {
    return pairGenerator<D, R, S>(this, T, cluster, s, rng);
  }

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

template <class U, int D, class R, class S> class SpinFlip : public Transformation<U, D, R, S> {
private:
  Spin<U, D, S>* sOld;
  Spin<U, D, S> sNew;

public:
  SpinFlip(Model<U, D, R, S>& M, const R& r, Spin<U, D, S>* s)
      : Transformation<U, D, R, S>(M, 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 = {NULL};
    std::set<Spin<U, D, S>*> oldNeighbors = this->M.dict.neighbors(sOld->x);
    std::set<Spin<U, D, S>*> newNeighbors = this->M.dict.neighbors(sNew.x);

    neighbors.insert(oldNeighbors.begin(), oldNeighbors.end());
    neighbors.insert(newNeighbors.begin(), newNeighbors.end());

    neighbors.erase(sOld);

    return neighbors;
  }

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

  Transformation<U, D, R, S>* createNew(double T, std::set<Spin<U, D, S>*>& cluster,
                                             Spin<U, D, S>* s, Rng& rng) const override {
    if (s == NULL) {
      return new FieldFlip<U, D, R, S>(this->M, this->r);
    } else {
      return pairGenerator<D, R, S>(this, T, cluster, s, rng);
    }
  }

  void apply() override {
    this->M.dict.remove(sOld);
    *sOld = sNew;
    this->M.dict.insert(sOld);
  }
};

template <class U, int D, class R, class S> class PairFlip : public Transformation<U, D, R, S> {
private:
  Spin<U, D, S>* s1Old;
  Spin<U, D, S>* s2Old;
  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)
      : Transformation<U, D, R, S>(M, 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 = {NULL};

    std::set<Spin<U, D, S>*> current_neighbors_1 = this->M.dict.neighbors(s1Old->x);
    std::set<Spin<U, D, S>*> current_neighbors_2 = this->M.dict.neighbors(s2Old->x);
    std::set<Spin<U, D, S>*> new_neighbors_1 = this->M.dict.neighbors(s1New.x);
    std::set<Spin<U, D, S>*> new_neighbors_2 = this->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.erase(s1Old);
    neighbors.erase(s2Old);

    return neighbors;
  }

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

  Transformation<U, D, R, S>* createNew(double T, std::set<Spin<U, D, S>*>& cluster,
                                             Spin<U, D, S>* s, Rng& rng) const override {
    if (s == NULL) {
      return new FieldFlip<U, D, R, S>(this->M, this->r);
    } else {
      return pairGenerator<D, R, S>(this, T, cluster, s, rng);
    }
  }

  void apply() override {
    this->M.dict.remove(s1Old);
    this->M.dict.remove(s2Old);
    *s1Old = s1New;
    *s2Old = s2New;
    this->M.dict.insert(s1Old);
    this->M.dict.insert(s2Old);
  }
};


template <int D, class R, class S>
Transformation<double, D, R, S>*
pairGenerator(const Transformation<double, D, R, S>* tOld, double T,
                        std::set<Spin<double, D, S>*>& cluster, Spin<double, D, S>* s, Rng& rng) {
  SpinFlip<double, D, R, S>* t = new SpinFlip<double, D, R, S>(tOld->M, tOld->r, s);
  Spin<double, D, S>* sMax = NULL;
  double ΔEMax = 0.0;
  for (Spin<double, D, S>* ss : t->toConsider()) {
    if (ss != NULL && ss != s && !cluster.contains(ss)) {
      double ΔE = t->ΔE(ss);
      if (ΔE > ΔEMax) {
        sMax = ss;
        ΔEMax = ΔE;
      }
    }
  }
  if (sMax == NULL || 1.0 - exp(-ΔEMax / T) < rng.uniform<double>(0, 1) || cluster.contains(sMax)) {
    return t;
  } else {
    delete t;
    cluster.insert(sMax);
    return new PairFlip<double, D, R, S>(tOld->M, tOld->r, s, sMax);
  }
}

template <int D, class R, class S>
Transformation<signed, D, R, S>*
pairGenerator(const Transformation<signed, D, R, S>* tOld, double T,
                        std::set<Spin<signed, D, S>*>& cluster, Spin<signed, D, S>* s, Rng& rng) {
  Vector<signed, 2> v = tOld->r.act(*s).x;
  std::set<Spin<signed, D, S>*> on_site = tOld->M.dict.at(v);
  if (on_site.empty()) {
    return new SpinFlip<signed, D, R, S>(tOld->M, tOld->r, s);
  } else {
    cluster.insert(*on_site.begin());
    return new PairFlip<signed, D, R, S>(tOld->M, tOld->r, s, *on_site.begin());
  }
}