#include "space_wolff.hpp"
#include <GL/glut.h>

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

class animation : public measurement<double, 2, Euclidean<double, D>, double> {
  public:
    animation(double L, unsigned w, int argc, char *argv[]) {
      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 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();
      getchar();
    }
};

std::function<Euclidean<double, D>(const model&, randutils::mt19937_rng&)> eGen(const std::vector<Matrix<double, 2>>& mats, const std::vector<Vector<double, 2>>& vecs) {
  return [&mats, &vecs] (const model& M, randutils::mt19937_rng& rng) -> Euclidean<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) += rng.variate<double, std::normal_distribution>(0.0, 0.1);
      }
      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()];
    }

    Euclidean<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);
    }
  }

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

      double rad_sum = pow(s1.s + s2.s, 2);

      bool overlap = d.transpose() * d < rad_sum;

      if (overlap) {
        return -1e8;
      } 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);
  animation A(L, 750, argc, argv);
  model sphere(L, Z, B, g, std::floor(log2(L)), 2, A);

  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 + rng.uniform(0.0, L / (4 * nx)), (i % nx) * L / nx + rng.uniform(0.0, L / (4 * nx))};
    sphere.s.push_back({pos, 0.5});
    sphere.dict.insert(&sphere.s.back());
  }


  sphere.wolff(T, N);

  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;
}