#include <fstream>
#include <iostream>
#include <chrono>

#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> {
private:
  uint64_t t1;
  uint64_t t2;
  unsigned n;
  unsigned tmp;
  unsigned wait;
  double L;
  Euclidean<double, D> s0_tmp;
  std::ofstream& file;

public:
  animation(double L, unsigned w, int argc, char* argv[], unsigned wait, std::ofstream& file) : s0_tmp(0), wait(wait), L(50 * L), file(file) {
    t1 = 0;
    t2 = 0;
    n = 0;

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

  void pre_cluster(const model& m, unsigned, const Transformation<double, D, Euclidean<double, D>, double>* t) override {
    s0_tmp = m.s0.inverse();
    tmp = 0;
    glClearColor(1.0f, 1.0f, 1.0f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);
    glColor3d(1.0, 0.0, 0.0);
    glBegin(GL_LINES);
      Vector<double, 2> r = t->r.t / 2;
      double θ = atan2(r(1), r(0));
      Vector<double, 2> v1 = s0_tmp.act({r(0) + L * sin(θ), r(1) - L * cos(θ)});
      Vector<double, 2> v2 = s0_tmp.act({r(0) - L * sin(θ), r(1) + L * cos(θ)});
      glVertex2d(v1(0), v1(1));
      glVertex2d(v2(0), v2(1));
      if (n > wait)
        file << v1.transpose() << " " << v2.transpose() << std::endl;
    glEnd();
    for (const Spin<double, 2, double>* s : m.s) {
      glBegin(GL_POLYGON);
      unsigned n_points = 50;
      glColor3f(0.0f, 0.0f, 0.0f);
      if (n > wait)
        file << s << " " << s->s << " " << s0_tmp.act(*s).x.transpose() << " ";
      for (unsigned i = 0; i < n_points; i++) {
        glVertex2d(s0_tmp.act(*s).x(0) + s->s * cos(2 * i * M_PI / n_points),
                   s0_tmp.act(*s).x(1) + s->s * sin(2 * i * M_PI / n_points));
      }
      glEnd();
    }
    if (n > wait)
      file << std::endl;
    glLineWidth(3);
    glFlush();
  }

  void plain_site_transformed(const model& m, const Transformation<double, D, Euclidean<double, D>, double>& t) override {
    if (n > wait) {
      if (t.current().size() > 1) {
        glColor3f(0.0f, 0.0f, 1.0f);
      } else {
        glColor3f(0.0f, 1.0f, 0.0f);
      }
      for (const Spin<double, 2, double>* s : t.current()) {
        if (s != NULL) {
          file << s << " " << s0_tmp.act(t.r.act(*s)).x.transpose() << " ";
          glBegin(GL_POLYGON);
          unsigned n_points = 50;
          for (unsigned i = 0; i < n_points; i++) {
            glVertex2d(s0_tmp.act(*s).x(0) + s->s * cos(2 * i * M_PI / n_points),
                       s0_tmp.act(*s).x(1) + s->s * sin(2 * i * M_PI / n_points));
          }
          glEnd();
        } else {
          file << s << " " << s0_tmp.act({0, 0}).transpose() << " ";
        }
      }
    }
    tmp++;
  }

  void post_cluster(const model& m) override {
    t1 += tmp;
    t2 += tmp * tmp;
    if (n > wait) {
      file << std::endl;
      glFlush();
      sleep(2);
    }
    n++;
  }

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

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

Gen<double, D, Euclidean<double, D>, double> eGen(double ε) {
  return [ε](model& M, Rng& rng) -> Transformation<double, D, Euclidean<double, D>, double>* {
    Vector<double, D> t;
    Matrix<double, D> m;

    double θ = rng.uniform((double)0.0, 2 * M_PI);
    m(0, 0) = -cos(2 * θ);
    m(1, 1) = cos(2 * θ);
    m(0, 1) = -2 * cos(θ) * sin(θ);
    m(1, 0) = -2 * cos(θ) * sin(θ);

    Spin<double, D, double>* s = rng.pick(M.s);
    t = s->x;
    for (unsigned j = 0; j < D; j++) {
      t(j) += rng.variate<double, std::normal_distribution>(0.0, ε);
    }

    Euclidean<double, D> g(t - m * t, m);
    return new SpinFlip<double, D, Euclidean<double, D>, double>(M, g, s);
  };
}

Gen<double, D, Euclidean<double, D>, double> mGen(double ε) {
  return [ε](model& M, Rng& rng) -> Transformation<double, D, Euclidean<double, D>, double>* {
    Matrix<double, D> m;

    Spin<double, D, double>* s1 = rng.pick(M.s);
    Spin<double, D, double>* s2 = rng.pick(M.s);

    while (s1 == s2) {
      s2 = rng.pick(M.s);
    }

    Vector<double, D> t1 = s1->x;
    Vector<double, D> t2 = s2->x;
    Vector<double, D> t12 = t1 - t2;
    Vector<double, D> t = (t1 + t2) / 2;

    double θ = atan2(t12(1), t12(0)) + rng.variate<double, std::normal_distribution>(0.0, ε) / t12.norm();

    m(0, 0) = -cos(2 * θ);
    m(1, 1) = cos(2 * θ);
    m(0, 1) = -2 * cos(θ) * sin(θ);
    m(1, 0) = -2 * cos(θ) * sin(θ);

    Vector<double, D> t3 = t - m * t;

    if (t3(0) != t3(0)) {
      std::cout << t3 << "\n" << t << "\n" << m << "\n" << t12 << "\n" ;
      getchar();
    }

    Euclidean<double, D> g(t3, m);
    return new PairFlip<double, D, Euclidean<double, D>, double>(M, g, s1, s2);
  };
}

Gen<double, D, Euclidean<double, D>, double> tGen(double ε) {
  return [ε](model& M, Rng& rng) -> Transformation<double, D, Euclidean<double, D>, double>* {
    Matrix<double, D> m;

    Spin<double, D, double>* s1 = rng.pick(M.s);
    Spin<double, D, double>* s2 = rng.pick(M.s);

    while (s1 == s2) {
      s2 = rng.pick(M.s);
    }

    Vector<double, D> t1 = s1->x;
    Vector<double, D> t2 = s2->x;
    Vector<double, D> t12 = t1 - t2;
    Vector<double, D> t = t2;

    double θ = atan2(t12(1), t12(0)) + rng.variate<double, std::normal_distribution>(0.0, ε) / t12.norm();

    m(0, 0) = -cos(2 * θ);
    m(1, 1) = cos(2 * θ);
    m(0, 1) = -2 * cos(θ) * sin(θ);
    m(1, 0) = -2 * cos(θ) * sin(θ);

    Vector<double, D> t3 = t - m * t;

    Euclidean<double, D> g(t3, m);
    return new SpinFlip<double, D, Euclidean<double, D>, double>(M, g, s1);
  };
}

Gen<double, D, Euclidean<double, D>, double> rGen(double ε) {
  return [ε](model& M, Rng& rng) -> Transformation<double, D, Euclidean<double, D>, double>* {
    Vector<double, D> t;
    Matrix<double, D> m;

    double θ = rng.uniform((double)0.0, 2 * M_PI);
    m(0, 0) = -cos(2 * θ);
    m(1, 1) = cos(2 * θ);
    m(0, 1) = -2 * cos(θ) * sin(θ);
    m(1, 0) = -2 * cos(θ) * sin(θ);

    Spin<double, D, double>* s = rng.pick(M.s);
    t = M.s0.t;
    for (unsigned j = 0; j < D; j++) {
      t(j) += rng.variate<double, std::normal_distribution>(0.0, ε);
    }

    Euclidean<double, D> g(t - m * t, m);
    return new SpinFlip<double, D, Euclidean<double, D>, double>(M, g, s);
  };
}

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;
  unsigned wait = 1000;

  double k = 1e2;
  double a = 0.1;

  int opt;

  while ((opt = getopt(argc, argv, "n:N:L:T:H:a:k:w:")) != -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;
    case 'a':
      a = atof(optarg);
      break;
    case 'k':
      k = atof(optarg);
      break;
    case 'w':
      wait = atoi(optarg);
      break;
    default:
      exit(1);
    }
  }

  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 = 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.norm();
  };

  std::function<double(Spin<double, D, double>)> B_hard = [L, H](Spin<double, D, double> s) -> double {
    if (fabs(s.x(0)) < 3 * L / 4 && fabs(s.x(1)) < 3 * L / 4) {
      return 0;
    } else {
      return std::numeric_limits<double>::infinity();
    }
  };

  auto g1 = eGen(1);
  auto g2 = mGen(0.1);
  auto g3 = tGen(0.1);
  auto g4 = rGen(0);

  auto tag = std::chrono::high_resolution_clock::now();

  std::string filename = "flips_" + std::to_string(n) + "_" + std::to_string(T) + "_" + std::to_string(H) + "_" + std::to_string(a) + "_" + std::to_string(k) + "_" +
                         std::to_string(tag.time_since_epoch().count()) + ".dat";

  std::ofstream file;
  file.open(filename);

  animation A(L, 750, argc, argv, wait, file);
  model sphere(1.0, Z, B);

  Rng rng;

  sphere.s.resize(n);

  unsigned nx = floor(sqrt(n));
  for (unsigned i = 0; i < sphere.s.size(); i++) {
    Spin<double, 2, double>* ss = new Spin<double, 2, double>();
    ss->x = {(i / nx) * L / nx - L / 2, (i % nx) * L / nx - L / 2};
    ss->s = rng.pick({0.45, 0.45});
    sphere.s[i] = ss;
    sphere.dict.insert(ss);
  }

  sphere.wolff(T, {g1, g2, g3, g4}, A, N);

  file.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.s << " " << rs.x.transpose() << "\n";
    delete s;
  }

  snapfile.close();

  return 0;
}