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

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

class animation : public measurement<signed, D, TorusGroup<signed, D>, signed> {
private:
  uint64_t t1;
  uint64_t t2;
  unsigned n;
  unsigned tmp;
  bool color;

public:
  animation(double L, unsigned w, bool tcolor, int argc, char* argv[]) {
    color = tcolor;
    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(0, L, 0, L);
  }

  void post_cluster(const model& m) override {
    glClearColor(1.0, 1.0, 1.0, 1.0);
    glClear(GL_COLOR_BUFFER_BIT);
    for (const Spin<signed, 2, signed>* s : m.s) {
      if (s->s == 1) {
        if (color)
          glColor3f(1.0, 0.0, 0.0);
        else
          glColor3f(1.0, 1.0, 1.0);
      } else if (s->s == -1) {
        if (color)
          glColor3f(0.0, 0.0, 1.0);
        else 
          glColor3f(0.0, 0.0, 0.0);
      }
      Vector<signed, 2> xx = m.s0.inverse().act(*s).x;
      glRecti(xx(0), xx(1), xx(0) + 1, xx(1) + 1);
    }
    glFlush();
  }
};

int main(int argc, char* argv[]) {

  unsigned L = 32;
  unsigned N = 1000;
  unsigned mod = 0;
  unsigned multi = 1e4;
  double mag = 0.5;
  double pop = 1.0;
  double T = 2.0 / log(1.0 + sqrt(2.0));
  double H = 1.0;
  double ε = 0.1;
  bool color = false;

  int opt;

  while ((opt = getopt(argc, argv, "N:L:T:H:e:m:M:r:p:c")) != -1) {
    switch (opt) {
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    case 'L':
      L = atoi(optarg);
      break;
    case 'T':
      T = atof(optarg);
      break;
    case 'H':
      H = atof(optarg);
      break;
    case 'e':
      ε = atof(optarg);
      break;
    case 'm':
      mod = atoi(optarg);
      break;
    case 'M':
      multi = atoi(optarg);
      break;
    case 'r':
      mag = atof(optarg);
      break;
    case 'p':
      pop = atof(optarg);
      break;
    case 'c':
      color = true;
      break;
    default:
      exit(1);
    }
  }

  double pZ = 1.0 - exp(-1.0 / T);

  std::function<double(const Spin<signed, D, signed>&, const Spin<signed, D, signed>&)> Z =
      [L, pZ](const Spin<signed, D, signed>& s1, const Spin<signed, D, signed>& s2) -> double {
    bool old_one_one = false;
    bool old_many_ones = false;
    bool old_any_two = false;

    Vector<signed, D> old_diff = diff<signed, D>(L, s1.x, s2.x);

    for (unsigned i = 0; i < D; i++) {
      if (old_diff(i) == 1 && !old_one_one) {
        old_one_one = true;
      } else if (old_diff(i) == 1 && old_one_one) {
        old_many_ones = true;
      } else if (old_diff(i) > 1) {
        old_any_two = true;
      }
    }

    bool were_on_someone = !old_one_one && !old_any_two;
    bool were_nearest_neighbors = old_one_one && !old_many_ones && !old_any_two;

    if (were_on_someone) {
      return -std::numeric_limits<double>::infinity();
      ;
    } else if (were_nearest_neighbors) {
      return s1.s * s2.s;
    } else {
      return 0.0;
    }
  };

  std::function<double(Spin<signed, D, signed>)> B;

  if (mod == 0) {
    B = [L, H](const Spin<signed, D, signed>& s) -> double {
      return H * s.s * smiley[15 - s.x(1) * 16 / L][s.x(0) * 16 / L];
    };
  } else {
    B = [mod, L, H](const Spin<signed, D, signed>& s) -> double {
      return H * s.s * sin(s.x(0) * mod * 2 * M_PI / L);
    };
  }


  Model<signed, D, TorusGroup<signed, D>, signed> ising(L, Z, B);

  Rng rng;

  unsigned n = 0;
  unsigned up = 0;
  unsigned down = 0;
  for (unsigned i = 0; i < L; i++) {
    for (unsigned j = 0; j < L; j++) {
      if (rng.uniform<double>(0, 1) < pop) {
      Spin<signed, D, signed>* ss = new Spin<signed, D, signed>();
      ss->x = {i, j};
      if (rng.uniform<double>(0, 1) < mag) {
        ss->s = 1;
        up++;
      } else {
        ss->s = -1;
        down++;
      }
      ising.s.push_back(ss);
      ising.dict.insert(ss);
      n++;
      }
    }
  }

  std::vector<Vector<signed, 2>> torusVectors = torus_vecs<signed, 2>(L);
  std::vector<Matrix<signed, 2>> torusMatrices = torus_mats<signed, 2>();

  Gen<signed, D, TorusGroup<signed, D>, signed> g =
      [L, &torusVectors,
       &torusMatrices](Model<signed, D, TorusGroup<signed, D>, signed>& M,
                       Rng& r) -> Transformation<signed, D, TorusGroup<signed, D>, signed>* {
    Matrix<signed, 2> m;
    Vector<signed, 2> t;

    m = r.pick(torusMatrices);
    t(0) = r.uniform<signed>(0, L);
    t(1) = r.uniform<signed>(0, L);
    t = t - m * t;

    TorusGroup<signed, 2> g = TorusGroup<signed, 2>({(signed)L, t, m});

    Spin<signed, 2, signed>* ss = r.pick(M.s);
    Spin<signed, 2, signed> s2 = g.act(*ss);

    while (ss->x(0) == s2.x(0) && ss->x(1) == s2.x(1)) {
      ss = r.pick(M.s);
      s2 = g.act(*ss);
    }

    std::set<Spin<signed, 2, signed>*> s2s = M.dict.at(s2.x);

    if (s2s.empty()) {
      return new SpinFlip<signed, 2, TorusGroup<signed, 2>, signed>(M, g, ss);
    } else {
      return new PairFlip<signed, 2, TorusGroup<signed, 2>, signed>(M, g, ss, *s2s.begin());
    }
  };

  animation A(L, 750, color, argc, argv);

  ising.wolff(T, {g}, A, N);

  /*
  for (unsigned i = 0; i < L; i++) {
    for (unsigned j = 0; j < L; j++) {
      if (i < L / 2) {
        ising.s.push_back({{i, j}, 1});
      } else {
        ising.s.push_back({{i, j}, -1});
      }
      ising.dict.record<signed>({i, j}, n);
      n++;
    }
  }
  */

  /*
  while (true) {
    ising.wolff(T, N, rng);
    std::array<double, 2> τ = ising.Eq.τ();
    std::cout << ising.Eq.num_added() << " " << τ[0] << " " << τ[1] << " " << τ[1] / τ[0] << "\n";
    if (τ[1] / τ[0] < ε && τ[0] * multi < ising.Eq.num_added()) {
      break;
    }
  }

  std::ofstream outfile;
  outfile.open("out.dat", std::ios::app);

  std::array<double, 2> act = ising.Eq.τ();
  std::vector<double> ρ = ising.Eq.ρ();

  outfile << L << " " << T << " " << mod << " " << H << " " << ising.Eq.num_added() << " " <<
  ising.Cq.avg() << " " << ising.Cq.serr() << " " << act[0] << " " << act[1]; for (double ρi : ρ) {
    outfile << " " << ρi;
  }
  outfile << "\n";

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

  std::vector<std::vector<signed>> snap(L);
  for (std::vector<signed>& line : snap) {
    line.resize(L);
  }

  for (spin<signed, D, signed> s : ising.s) {
    spin<signed, D, signed> snew = ising.s0.inverse().act(s);
    snap[snew.x(0)][snew.x(1)] = snew.s;
  }

  for (std::vector<signed> row : snap) {
    for (signed s : row) {
      snapfile << s << " ";
    }
    snapfile << "\n";
  }

  snapfile.close();
  */

  return 0;
}