#include <getopt.h>
#include <iostream>

#include "../include/randutils/randutils.hpp"

#include <functional>
#include <vector>
#include <queue>

#include <wolff/types.h>
#include <wolff/graph.hpp>

using namespace wolff;

namespace ising {
#define WOLFF_SITE_DEPENDENCE
#include <wolff.hpp>
#include <wolff/models/ising.hpp>
#undef WOLFF_SITE_DEPEDENCE
}

namespace xy {
#undef WOLFF_H
#undef WOLFF_SYSTEM_H
#undef WOLFF_CLUSTER_H
#undef WOLFF_MEASUREMENTS_H

#define WOLFF_NO_FIELD
#define WOLFF_BOND_DEPENDENCE
#include <wolff.hpp>
#include <wolff/models/vector.hpp>
#include <wolff/models/orthogonal.hpp>
#undef WOLFF_NO_FIELD
#undef WOLFF_BOND_DEPENDENCE
}

double E;

using ising::wolff::ising_t;
using xy::wolff::vector_t;
using xy::wolff::orthogonal_t;

class i_measurement : public ising::wolff::measurement<ising_t, ising_t> {
  private:
    N_t nc;
    N_t ns;

    int M;
    int A;
    v_t C;

    double totalaA;
    double totalA2;
    double totalaM;
    double totalM2;
    double totalC;
    double totalC2;

  public:
    i_measurement(const ising::wolff::system<ising_t, ising_t>& S) {
      nc = 0;
      ns = 0;
      M = 0;
      A = S.nv;

      totalaA = 0;
      totalA2 = 0;
      totalaM = 0;
      totalM2 = 0;
      totalC = 0;
      totalC2 = 0;
    }

    void pre_cluster(N_t, N_t, const ising::wolff::system<ising_t, ising_t>&, v_t, const ising_t&) override {
      C = 0;
    }

    void plain_bond_visited(const ising::wolff::system<ising_t, ising_t>&, v_t, const ising_t&, v_t, double dE) override {
      E += dE;
    }

    void ghost_bond_visited(const ising::wolff::system<ising_t, ising_t>& S, v_t i, const ising_t& s_old, const ising_t& s_new, double dE) override {
      E += dE;
      M += s_new - s_old;

      if (i >= S.nv / 2) {
        A -= s_new - s_old;
      } else {
        A += s_new - s_old;
      }
    }

    void plain_site_transformed(const ising::wolff::system<ising_t, ising_t>& S, v_t i, const ising_t& si_new) override {
      C++;
    }

    void post_cluster(N_t, N_t, const ising::wolff::system<ising_t, ising_t>&) override {
      totalaA += abs(A);
      totalA2 += pow(A, 2);
      totalaM += abs(M);
      totalM2 += pow(M, 2);
      totalC += C;
      totalC2 += pow(C, 2);

      nc++;

      ns++;
    }

    double avgaA() {
      return totalaA / ns;
    }

    double avgA2() {
      return totalA2 / ns;
    }

    double avgaM() {
      return totalaM / ns;
    }

    double avgM2() {
      return totalM2 / ns;
    }

    double avgC() {
      return totalC / nc;
    }

    double avgC2() {
      return totalC2 / nc;
    }
};

class x_measurement : public xy::wolff::measurement<orthogonal_t<2, double>, vector_t<2, double>> {
  private:
    N_t ns;
    N_t nc;
    orthogonal_t<2, double> R;

    vector_t<2, double> M;
    v_t C;
    vector_t<2, double> CM;

    vector_t<2, double> totalaM;
    vector_t<2, double> totalM2;
    double totalM;
    double totalCM;
    double totalC;
    double totalC2;

  public:
    x_measurement(const xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>>& S) {
      ns = 0;
      nc = 0;
      M = S.nv * S.s[0];

      totalaM.fill(0);
      totalM2.fill(0);
      totalC = 0;
      totalC2 = 0;
      totalM = 0;
      totalCM = 0;
    }

    void pre_cluster(N_t, N_t, const xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>>&, v_t, const orthogonal_t<2, double>& r) override {
      R = r;
      C = 0;
      CM.fill(0);
    }

    void plain_bond_visited(const xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>>&, v_t, const vector_t<2, double>&, v_t, double dE) override {
      E += dE;
    }

    void plain_site_transformed(const xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>>& S, v_t i, const vector_t<2, double>& si_new) override {
      C++;
      CM += S.s[i];
      M += si_new - S.s[i];
    }

    void post_cluster(N_t, N_t, const xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>>&) override {
      totalaM[0] += fabs(M[0]);
      totalaM[1] += fabs(M[1]);
      totalM2[0] += pow(M[0], 2); 
      totalM2[1] += pow(M[1], 2); 
      totalM += (pow(M[0], 2) + pow(M[1], 2));
      totalCM += pow(CM[0] * R[0][0] + CM[1] * R[0][1], 2) / C;
      totalC += C;
      totalC2 += pow(C, 2);

      nc++;

      ns++;
    }

    vector_t<2, double> avgaM() {
      return totalaM / ns;
    }

    vector_t<2, double> avgM2() {
      return totalM2 / ns;
    }

    double avgM() {
      return totalM / ns;
    }

    double avgCM() {
      return totalCM / ns;
    }

    double avgC() {
      return totalC / nc;
    }

    double avgC2() {
      return totalC2 / nc;
    }
};

using namespace wolff;

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

  // set defaults
  N_t N = (N_t)1e4;
  const D_t D = 2;
  L_t L = 128;
  double T = 2.26918531421;
  double J = 1.0;
  double a = 1.5;

  int opt;

  // take command line arguments
  while ((opt = getopt(argc, argv, "N:L:T:J:a:")) != -1) {
    switch (opt) {
    case 'N': // number of steps
      N = (N_t)atof(optarg);
      break;
    case 'L': // linear size
      L = atoi(optarg);
      break;
    case 'T': // temperature
      T = atof(optarg);
      break;
    case 'J': // temperature
      J = atof(optarg);
      break;
    case 'a': // temperature
      a = atof(optarg);
      break;
    default:
      exit(EXIT_FAILURE);
    }
  }

  randutils::auto_seed_128 seeds;
  std::mt19937 rng{seeds};

  std::function <double(const ising_t&, const ising_t&)> Zi = [=] (const ising_t& s1, const ising_t s2) -> double {
    if (s1.x == s2.x) {
      return -J;
    } else {
      return J;
    }
  };

  std::vector<vector_t<2, double>> *xy_spins;

  std::function <double(v_t, const ising_t&)> Bi = [L, D, a, &xy_spins] (v_t v, const ising_t& s) -> double {
    v_t xyv1 = v % (v_t)pow(L, D);
    v_t xyv2;
    if (v / (v_t)pow(L, D) == 0) {
      xyv2 = L * (xyv1 / L) + ((xyv1 % L) + 1) % L;
    } else {
      xyv2 = L * (((xyv1 / L) + 1) % L) + (xyv1 % L);
    }

    double value = a * ((xy_spins->at(xyv1)) * (xy_spins->at(xyv2)));

    if (s.x) {
      return -value;
    } else {
      return value;
    }
  };

  graph Gi;
  
  Gi.nv = 2 * pow(L, 2);
  Gi.ne = 2 * Gi.nv;

  Gi.adjacency.resize(Gi.nv);
  Gi.coordinate.resize(Gi.nv);

  for (std::vector<v_t> adj_i : Gi.adjacency) {
    adj_i.reserve(4);
  }

  for (D_t i = 0; i < 2; i++) {
    v_t sb = i * pow(L, 2);

    for (v_t j = 0; j < pow(L, 2); j++) {
      v_t vc = sb + j;

      Gi.adjacency[vc].push_back(((i + 1) % 2) * pow(L, 2) + j);
      Gi.adjacency[vc].push_back(((i + 1) % 2) * pow(L, 2) + pow(L, 2 - 1) * (j / (v_t)pow(L, 2 - 1)) + (j + 1 - 2 * (i % 2)) % L);
      Gi.adjacency[vc].push_back(((i + 1) % 2) * pow(L, 2) + pow(L, 2 - 1) * ((L + (j/ (v_t)pow(L, 2 - 1)) - 1 + 2 * (i % 2)) % L) + (j - i) % L);
      Gi.adjacency[vc].push_back(((i + 1) % 2) * pow(L, 2) + pow(L, 2 - 1) * ((L + (j/ (v_t)pow(L, 2 - 1)) - 1 + 2 * (i % 2)) % L) + (j + 1 - i) % L);
    }
  }
  
  ising::wolff::system<ising_t, ising_t> si(Gi, T, Zi, Bi);

  std::function <double(v_t, const vector_t<2, double>&, v_t, const vector_t<2, double>&)> Zxy = [L, D, a, &si] (v_t v1, const vector_t<2, double>& s1, v_t v2, const vector_t<2, double>& s2) -> double {
    v_t iv;
    if (v1 / L == v2 / L) {
      v_t vs = v1 < v2 ? v1 : v2;
      v_t vl = v1 < v2 ? v2 : v1;

      if (vl == L - 1 && vs == 0) {
        iv = vl;
      } else {
        iv = vs;
      }
    } else {
      v_t vs = v1 < v2 ? v1 : v2;
      v_t vl = v1 < v2 ? v2 : v1;

      if (vl / L == L - 1 && vs / L == 0) {
        iv = vl;
      } else {
        iv = vs;
      }
    }

    if (si.s[iv].x) {
      return (1 + a) * (s1 * s2);
    } else {
      return (1 - a) * (s1 * s2);
    }
  };

  graph Gxy(2, L);

  xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>> sxy(Gxy, T, Zxy);

  xy_spins = &(sxy.s);

  for (v_t i = 0; i < pow(L, D); i++) {
    si.s[pow(L, D) + i].x = true;
  }

  E = - J * 4 * pow(L, D) - 2 * pow(L, D);

  i_measurement mi(si);
  x_measurement mxy(sxy);

  double total_E;
  double total_E2;

  for (N_t i = 0; i < N; i++) {
    si.run_wolff(1, ising::wolff::gen_ising, mi, rng);
    sxy.run_wolff(1, xy::wolff::generate_rotation_uniform<2>, mxy, rng);
    total_E += E;
    total_E2 += pow(E, 2);
  }

  std::cout << mi.avgC() / si.nv << " " << mi.avgA2() / pow(si.nv, 2) << "\n";
  std::cout << mxy.avgM() / pow(sxy.nv, 1) << " " << 2 * mxy.avgCM() / pow(sxy.nv, 0) << "\n";
}