#include <GL/glut.h>
#include <fstream>
#include <iostream>

#include "space_wolff.hpp"
#include "torus_symmetries.hpp"

const unsigned D = 2;
typedef Model<double, D, TorusGroup<double, D>, double> model;

class animation : public measurement<double, 2, TorusGroup<double, D>, double> {
private:
  uint64_t t1;
  uint64_t t2;
  unsigned n;
  unsigned tmp;

public:
  animation(double L, unsigned w, int argc, char* argv[]) {
    t1 = 0;
    t2 = 0;
    n = 0;

    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_SINGLE | GLUT_RGB);
    glutInitWindowSize(w, w);
    glutCreateWindow("wolff");
    glClearColor(0.0, 0.0, 0.0, 0.0);
    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    gluOrtho2D(-1, L + 1, -1, L + 1);
  }

  void pre_cluster(const model&, unsigned, const TorusGroup<double, D>&) override { tmp = 0; }

  void plain_site_transformed(const model&, const Spin<double, D, double>*,
                              const Spin<double, D, double>&) override {
    tmp++;
  }

  void post_cluster(const model& m) override {
    glClearColor(1.0f, 1.0f, 1.0f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);
    for (const Spin<double, 2, double>& s : m.s) {
      glBegin(GL_POLYGON);
      unsigned n_points = 50;
      glColor3f(0.0f, 0.0f, 0.0f);
      for (unsigned i = 0; i < n_points; i++) {
        glVertex2d(m.s0.inverse().act(s).x(0) + s.s * cos(2 * i * M_PI / n_points),
                   m.s0.inverse().act(s).x(1) + s.s * sin(2 * i * M_PI / n_points));
      }
      glEnd();
    }
    glFlush();

    t1 += tmp;
    t2 += tmp * tmp;
    n++;
  }

  void clear() {
    t1 = 0;
    t2 = 0;
    n = 0;
  }

  double var() { return (t2 - t1 * t1 / (double)n) / (double)n; }
};

std::function<TorusGroup<double, D>(const model&, randutils::mt19937_rng&)>
eGen(const std::vector<Matrix<double, 2>>& mats, const std::vector<Vector<double, 2>>& vecs,
     double ε, double L) {
  return
      [&mats, &vecs, L, ε](const model& M, randutils::mt19937_rng& rng) -> TorusGroup<double, 2> {
        Vector<double, D> t;
        Matrix<double, D> m;
        unsigned flip = rng.uniform((unsigned)0, (unsigned)(mats.size() + vecs.size() - 1));
        if (flip < mats.size()) {
          unsigned f_ind = rng.uniform((unsigned)0, (unsigned)M.s.size());
          t = M.s[f_ind].x;
          for (unsigned j = 0; j < D; j++) {
            t(j) = fmod(10 * L + t(j) + rng.variate<double, std::normal_distribution>(0.0, ε), L);
          }
          m = mats[flip];
        } else {
          for (unsigned j = 0; j < D; j++) {
            for (unsigned k = 0; k < D; k++) {
              if (j == k) {
                m(j, k) = 1;
              } else {
                m(j, k) = 0;
              }
            }
          }

          t = vecs[flip - mats.size()];
        }

        TorusGroup<double, D> g(M.L, t, m);
        return g;
      };
}

int main(int argc, char* argv[]) {
  const unsigned D = 2;

  double L = 32;
  unsigned N = 1000;
  double T = 2.0 / log(1.0 + sqrt(2.0));
  double H = 1.0;
  unsigned n = 25;

  int opt;

  while ((opt = getopt(argc, argv, "n:N:L:T:H:")) != -1) {
    switch (opt) {
    case 'n':
      n = (unsigned)atof(optarg);
      break;
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    case 'L':
      L = atof(optarg);
      break;
    case 'T':
      T = atof(optarg);
      break;
    case 'H':
      H = atof(optarg);
      break;
    default:
      exit(1);
    }
  }

  double k = 1e2;
  double a = 0.05;

  std::function<double(const Spin<double, D, double>&, const Spin<double, D, double>&)> Z =
      [L, a, k](const Spin<double, D, double>& s1, const Spin<double, D, double>& s2) -> double {
    Vector<double, D> d = diff(L, s1.x, s2.x);

    double σ = s1.s + s2.s;
    double δ = σ - sqrt(d.transpose() * d);

    if (δ > -a * σ) {
      return 0.5 * k * (2 * pow(a * σ, 2) - pow(δ, 2));
    } else if (δ > -2 * a * σ) {
      return 0.5 * k * pow(δ + 2 * a * σ, 2);
    } else {
      return 0;
    }
  };

  std::function<double(Spin<double, D, double>)> B = [L, H](Spin<double, D, double> s) -> double {
    return H * s.x(1);
  };

  std::vector<Matrix<double, D>> mats = torus_mats<double, D>();
  std::vector<Vector<double, D>> vecs = torus_vecs<double, D>(L);
  auto g = eGen(mats, vecs, L, L);
  animation A(L, 750, argc, argv);
  model sphere(L, Z, B);

  randutils::mt19937_rng rng;

  sphere.s.reserve(n);

  unsigned nx = floor(sqrt(n));
  for (unsigned i = 0; i < n; i++) {
    Vector<double, D> pos = {(i / nx) * L / nx, (i % nx) * L / nx};
    sphere.s.push_back({pos, 0.5});
    sphere.dict.insert(&sphere.s.back());
  }

  sphere.wolff(T, g, A, N);
  std::ofstream outfile;
  outfile.open("test.dat");

  for (signed i = -10; i <= 10; i++) {
    A.clear();
    double ε = pow(2, -4 + i / 2.0);
    auto gn = eGen(mats, vecs, ε, L);
    sphere.wolff(T, gn, A, N);
    outfile << ε << " " << A.var() / sphere.s.size() << std::endl;
    std::cout << ε << " " << A.var() / sphere.s.size() << std::endl;
  }

  outfile.close();

  std::ofstream snapfile;
  snapfile.open("sphere_snap.dat");

  for (Spin<double, D, double> s : sphere.s) {
    Spin<double, D, double> rs = sphere.s0.inverse().act(s);
    snapfile << rs.x.transpose() << "\n";
  }

  snapfile.close();

  return 0;
}