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

#include "glass.hpp"

template <unsigned D> class CiamarraState : public Vector<D> {
public:
  CiamarraState() : Vector<D>(Vector<D>::Zero()) {}

  CiamarraState(const Vector<D>& v) : Vector<D>(v) {}

  CiamarraState(unsigned a, signed b) : Vector<D>(Vector<D>::Zero()) { CiamarraState::operator()(a) = b; }

  CiamarraState(Rng& r) : CiamarraState(r.uniform((unsigned)0, D - 1), r.pick({-1, 1})) {}

  bool empty() const { return CiamarraState::squaredNorm() == 0; }

  void remove() { CiamarraState::setZero(); }

  bool operator==(const CiamarraState<D>& s) const {
    return CiamarraState::dot(s) == 1;
  }

  CiamarraState<D> flip() const {
    CiamarraState<D> s;
    for (unsigned i = 0; i < D; i++) {
      s(i) = -this->operator()(i);
    }
    return s;
  }

  CiamarraState<D> transform(const Transformation<D>& m) const {
    return m * *this;
  }
};

template <unsigned D>
class CiamarraSystem : public HardSystem<D, CiamarraState<D>> {
  public:
  using HardSystem<D, CiamarraState<D>>::HardSystem;

  std::list<std::reference_wrapper<Vertex<D, CiamarraState<D>>>>
    overlaps(Vertex<D, CiamarraState<D>>& v, const CiamarraState<D>& s, bool
        excludeSelf = false) override {
    std::list<std::reference_wrapper<Vertex<D, CiamarraState<D>>>> o;

    if (s.empty()) {
      return o;
    }

    if (!v.empty() && !excludeSelf) {
      o.push_back(v);
    }

    for (const HalfEdge<D, CiamarraState<D>>& e : v.adjacentEdges) {
      if (!e.neighbor.empty()) {
        if (s.dot(e.Δx) == 1 || e.neighbor.state.dot(e.Δx) == -1) {
          o.push_back(e.neighbor);
        }
      }
    }

    return o;
  }

  void setGroundState() {
    CiamarraSystem::N = 0;

    for (Vertex<D, CiamarraState<D>>& v : CiamarraSystem::vertices) {
      unsigned a = 0;
      for (unsigned d = 0; d < D; d++) {
        a += (d + 1) * v.position(d);
      }
      a %= 2 * D + 1;

      v.state.setZero() = Vector<D>::Zero();

      if (0 < a && a <= D) {
        v.state(a - 1) = -1;
        CiamarraSystem::N++;
      } else if (D < a) {
        v.state(2 * D - a) = 1;
        CiamarraSystem::N++;
      }
    }
  }

  /* For the Ciamarra system, position within sites must be specially acconted
   * for in the scattering function. We therefore expand the lattice and add
   * the state when setting positions, so as later to be able to take particle
   * sublattice positions into account.
   * */
  void setInitialPosition() override {
    CiamarraSystem::orientation = Transformation<D>(CiamarraSystem::L);
    for (Vertex<D, CiamarraState<D>>& v : CiamarraSystem::vertices) {
      v.initialPosition = 4 * v.position + v.state;
    }
  }

  double selfIntermediateScattering(const std::vector<Matrix<D>>& ms) const override {
    double F = 0;

    std::array<Vector<D>, D> ks;
    for (unsigned i = 0; i < D; i++) {
      ks[i].setZero();
      ks[i](i) = 1;
    }
    for (const Vertex<D, CiamarraState<D>>& v : CiamarraSystem::vertices) {
      if (!v.empty()) {
        Vector<D> Δx = ((4 * CiamarraSystem::orientation.inverse().apply(v.position)) + CiamarraSystem::orientation.inverse().apply(v.state)) - v.initialPosition;
        for (const Vector<D>& k : ks) {
          F += cos((M_PI / 4) * k.dot(Δx));
        }
      }
    }

    return F / (D * CiamarraSystem::N);
  }
};

void print(const CiamarraSystem<2>& s) {
  for (const Vertex<2, CiamarraState<2>>& v : s.vertices) {
    if (v.state(0) == 1 && v.state(1) == 0) {
      std::cerr << "▶";
    } else if (v.state(0) == -1 && v.state(1) == 0) {
      std::cerr << "◀";
    } else if (v.state(0) == 0 && v.state(1) == -1) {
      std::cerr << "▲";
    } else if (v.state(0) == 0 && v.state(1) == 1) {
      std::cerr << "▼";
    } else if (v.state(0) == 0 && v.state(1) == 0) {
      std::cerr << " ";
    } else {
      std::cerr << "X";
    }

    if (v.position(0) == s.L - 1) {
      std::cerr << std::endl;
    }
  }
}

int main() {
  const unsigned D = 3;
  unsigned L = 15;
  unsigned Nmin = 2e2;
  unsigned Nmax = 2e5;
  double Tmin = 0.04;
  double Tmax = 0.2;
  double δT = 0.02;

  CiamarraSystem<D> s(L);

  Rng r;

  double z = exp(1 / 0.08);

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

  for (double ρ : {0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.76, 0.78, 0.8, 0.81, 0.815, 0.818, 0.821, 0.825}) {
    while (s.density() < ρ) {
      s.sweepGrandCanonical(z, r);
    }
    std::cout << s.density() << " ";
    s.setInitialPosition();
    auto start = std::chrono::high_resolution_clock::now();
    for (unsigned i = 0; i < 1e4; i++) {
      /*
      for (unsigned j = 0; j < s.vertices.size(); j++) {
        s.tryRandomNonlocalMove(r);
      }
      */
      Matrix<D> m = r.pick(ms);
      Vertex<D, CiamarraState<D>>& v = r.pick(s.vertices);
      unsigned n = s.wolff(Transformation<D>(L, m, v.position - m * v.position), r);
      if (i % 1 == 0) {
        auto stop = std::chrono::high_resolution_clock::now();
        auto duration = duration_cast<std::chrono::microseconds>(stop - start);
        std::cout << duration.count() << " " << s.selfIntermediateScattering(ms) << " " << n << " ";
      }
    }
    if (!s.compatible()) {
      std::cerr << "Bad configuration" << std::endl;
      return 1;
    }
    std::cout << std::endl;
  }

  return 0;
}