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

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

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

namespace ising {
#define WOLFF_SITE_DEPENDENCE
#include <wolff_models/ising.hpp> // ising includes wolff
#undef WOLFF_SITE_DEPEDENCE
  using namespace wolff;
}

namespace xy {
#undef WOLFF_H

#define WOLFF_NO_FIELD
#define WOLFF_BOND_DEPENDENCE
#include <wolff_models/vector.hpp>
#include <wolff_models/orthogonal.hpp> // orthogonal includes wolff
#undef WOLFF_NO_FIELD
#undef WOLFF_BOND_DEPENDENCE
  using namespace wolff;
}

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

// at each ising site, we need to know which sublattice we're in and the xy
// spins to our left and right (top and bottom)
typedef struct _ising_site_props {
  unsigned sublattice;
  vector_t<2, double>* back;
  vector_t<2, double>* front;
} ising_site_props;

// at each xy bond, we need to know whether we're nearest or next-nearest, the
// displacement the bond makes, and the ising spin that lies on us (if any)
typedef struct _xy_bond_props {
  unsigned distance;
  vector_t<2, int> dNeighborSelf;
  ising_t* spin;
} xy_bond_props;

typedef ising::graph<ising_site_props> graph_i;
typedef xy::graph<std::tuple<>, xy_bond_props> graph_x;

// the sin of the angle between two vectors, signed: vsin(v1, v2) = -vsin(v2, v1)
double vsin(const vector_t<2, double>& v1, const vector_t<2, double>& v2) {
  return v1[0] * v2[1] - v1[1] * v2[0];
}

// these need to be global so that both wolff hooks have access to them.
// TODO: investigate using pointers
double cos_xy;
vector_t<2, double> sin_xy;

std::vector<double> ρ(unsigned N, const std::vector<double>& C, double avg) {
  double C0 = C.front() / N;
  double avg2 = pow(avg / N, 2);

  std::vector<double> ρtmp;

  for (double Ct : C) {
    ρtmp.push_back((Ct / N - avg2) / (C0 - avg2));
  }

  return ρtmp;
}

double τ(unsigned n, const std::vector<double>& C, double avg) {
  double τtmp = 0.5;
  unsigned M = 1;
  double c = 8.0;

  std::vector<double> ρ_tmp = ρ(n, C, avg);

  while (c * τtmp > M && M < C.size()) {
    τtmp += ρ_tmp[M];
    M++;
  }

  return τtmp;
}

class quantity {
  private:
    double total;
    std::vector<double> C;

  public:
    unsigned n;
    std::list<double> hist;
    quantity(unsigned lag) : C(lag) {
      n = 0;
      total = 0;
    }

    void add(double x) {
      hist.push_front(x);
      if (hist.size() > C.size()) {
        hist.pop_back();
        unsigned t = 0;
        for (double a : hist) {
          C[t] += a * x;
          t++;
        }
        total += x;
        n++;
      }
    }

    double avg() {
      return total / n;
    }

    double σ() {
      return 2.0 / n * τ(n, C, total) * (C[0] / n - pow(this->avg(), 2));
    }

    double serr() {
      return sqrt(this->σ());
    }
};

class quantity_pair {
  private:
    std::vector<double> Cab;
    double total;
    double T;
    quantity a;
    quantity b;

  public:
    quantity_pair(unsigned lag, double T) : a(lag), b(lag), Cab(lag), T(T) {};

    void add(double x, double y) {
      a.add(x);
      b.add(y);

      if (a.hist.size() == Cab.size()) {
        unsigned t = 0;
        auto a_it = a.hist.begin();
        auto b_it = b.hist.begin();
        while (a_it != a.hist.end()) {
          Cab[t] += (*a_it) * (*b_it);
          t++;
          a_it++;
          b_it++;
        }
      }
    }

    double cov() {
      return 1.0 / a.n * τ(a.n, Cab, sqrt(a.avg() * b.avg()) * a.n) * (Cab[0] / a.n - a.avg() * b.avg());
    }

    double avg() {
      return a.avg() - b.avg() / T;
    }

    double serr() {
      return sqrt(a.σ() + b.σ() / pow(T, 2) - 2 * this->cov() / T);
    }
};

class i_measurement : public ising::wolff::measurement<ising_t, ising_t, graph_i> {
  private:
    double a;
    double K1;
    unsigned w;
    int A;
    unsigned N;

  public:
    quantity A2;

    i_measurement(const ising::wolff::system<ising_t, ising_t, graph_i>& S, double a, double K1, unsigned nh, unsigned w) : a(a), w(w), K1(K1), A2(nh) {
      A = S.nv;
      N = 0;
    }

    void ghost_bond_visited(const ising::wolff::system<ising_t, ising_t, graph_i>& S, const graph_i::vertex& v, const ising_t& s_old, const ising_t& s_new, double dE) override {
      if (v.prop.sublattice == 0) {
        A += s_new - s_old;
      } else {
        A -= s_new - s_old;
      }

      sin_xy[v.prop.sublattice] += K1 * a * (s_new - s_old) * vsin(*(v.prop.back), *(v.prop.front));

      cos_xy -= dE;
    }

    void post_cluster(unsigned, unsigned, const ising::wolff::system<ising_t, ising_t, graph_i>&) override {
      if (N > w) {
        A2.add(pow(A, 2));
      }
      N++;
    }
};

class x_measurement : public xy::wolff::measurement<orthogonal_t<2, double>, vector_t<2, double>, graph_x> {
  private:
    double K1y;
    double K1yn;
    double a;
    const ising_t* si0;

  public:
    x_measurement(const xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>, graph_x>& S, double K1y, double K1yn, double a, const ising_t *si0) : K1y(K1y), K1yn(K1yn), a(a), si0(si0) {
    }

    void plain_bond_visited(const xy::wolff::system<orthogonal_t<2, double>, vector_t<2, double>, graph_x>& S, const graph_x::halfedge& e, const vector_t<2, double>& s1_new, double dE) override {

      vector_t<2, double> dsin = (vsin(s1_new, S.s[e.neighbor.ind]) - vsin(S.s[e.self.ind], S.s[e.neighbor.ind])) * e.prop.dNeighborSelf;

      if (e.prop.distance == 0) {
        cos_xy -= dE;
        sin_xy += K1y * (1 + (*si0) * a * *(e.prop.spin)) * dsin;
      } else if (e.prop.distance == 1) {
        cos_xy -= 2 * dE;
        sin_xy += K1y * K1yn * dsin;
      }
    }
};

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

  // set defaults
  unsigned N = (unsigned)1e4; // number of steps between precision checks
  unsigned w = (unsigned)1e3; // steps to skip before data collection begins
  unsigned nh = 1e2; // lag to keep in autocorrelation functions
  double εmag = 1e-1; // precision goal in magnetization
  double εstiff = 1e-1; // precision goal in spin stiffness

  const unsigned D = 2; // we're always in 2D
  unsigned L = 128;     // system size
  double T = 1.0;       // temperature
  double J = 1.0;      // ising nn coupling
  double K1 = 1.0;      // xy nn coupling
  double K2 = 0.0;      // xy nnn coupling
  double α = 1.5;       // ising-xy coupling

  int opt;

  // take command line arguments
  while ((opt = getopt(argc, argv, "N:L:T:J:K:a:n:e:f:")) != -1) {
    switch (opt) {
      case 'N': // number of steps
        N = (unsigned)atof(optarg);
        break;
      case 'L': // linear size
        L = atoi(optarg);
        break;
      case 'T':
        T = atof(optarg);
        break;
      case 'J':
        J = atof(optarg);
        break;
      case 'K':
        K1 = atof(optarg);
        break;
      case 'n':
        K2 = atof(optarg);
        break;
      case 'a':
        α = atof(optarg);
        break;
      case 'e':
        εmag = atof(optarg);
        break;
      case 'f':
        εstiff = atof(optarg);
        break;
      default:
        exit(EXIT_FAILURE);
    }
  }

  // initialize random numbers
  randutils::auto_seed_128 seeds;
  std::mt19937 rng{seeds};

  // antiferromagnetic Ising coupling
  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;
    }
  };

  // ising field depends on dot product of surrounding xy spins
  std::function <double(const graph_i::vertex&, const ising_t&)> Bi =
    [K1, α] (const graph_i::vertex& v, const ising_t& s) -> double {
      return K1 * α * s * (*(v.prop.back) * *(v.prop.front));
    };

  // initialize Ising diagonal-lattice graph
  graph_i Gi;

  Gi.L = L;
  Gi.D = 2;

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

  Gi.vertices.resize(Gi.nv);

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

    for (unsigned j = 0; j < pow(L, 2); j++) {
      unsigned vc = sb + j;
      Gi.vertices[vc].ind = vc;
      Gi.vertices[vc].prop.sublattice = i;

      graph_i::halfedge e1(Gi.vertices[vc], Gi.vertices[((i + 1) % 2) * pow(L, 2) + j]);
      graph_i::halfedge e2(Gi.vertices[vc], Gi.vertices[((i + 1) % 2) * pow(L, 2) + pow(L, 2 - 1) * (j / (unsigned)pow(L, 2 - 1)) + (j + 1 - 2 * (i % 2)) % L]);
      graph_i::halfedge e3(Gi.vertices[vc], Gi.vertices[((i + 1) % 2) * pow(L, 2) + pow(L, 2 - 1) * ((L + (j/ (unsigned)pow(L, 2 - 1)) - 1 + 2 * (i % 2)) % L) + (j - i) % L]);
      graph_i::halfedge e4(Gi.vertices[vc], Gi.vertices[((i + 1) % 2) * pow(L, 2) + pow(L, 2 - 1) * ((L + (j/ (unsigned)pow(L, 2 - 1)) - 1 + 2 * (i % 2)) % L) + (j + 1 - i) % L]);

      Gi.vertices[vc].edges.push_back(e1);
      Gi.vertices[vc].edges.push_back(e2);
      Gi.vertices[vc].edges.push_back(e3);
      Gi.vertices[vc].edges.push_back(e4);
    }
  }

  ising::wolff::system<ising_t, ising_t, graph_i> si(Gi, T, Zi, Bi);

  // ferromagnetic XY coupling, with nearest-neighbor energy that deponds on
  // state of Ising spin on the bond
  std::function <double(const graph_x::halfedge&, const vector_t<2, double>&, const vector_t<2, double>&)>
    Zxy = [K1, K2, α, &si] (const graph_x::halfedge& e, const vector_t<2, double>& s1, const vector_t<2, double>& s2) -> double {
      if (e.prop.distance == 0) {
        return K1 * (1 + si.s0 * α * *(e.prop.spin)) * (s1 * s2);
      } else if (e.prop.distance == 1) {
        return K1 * K2 * (s1 * s2);
      } else {
        return 0;
      }
  };

  // initialize square-lattice xy graph
  graph_x Gxy(2, L);

  // assign edge properties to the XY nearest-neighbor bonds
  for (graph_x::vertex& v : Gxy.vertices) {
    for (graph_x::halfedge& e : v.edges) {
      e.prop.distance = 0;

      int v1 = e.self.ind;
      int v2 = e.neighbor.ind;

      e.prop.dNeighborSelf[0] = (v2 % L) - (v1 % L);
      e.prop.dNeighborSelf[1] = v2 / L - v1 / L;

      for (unsigned i = 0; i < 2; i++) {
        if (e.prop.dNeighborSelf[i] == L - 1) {
          e.prop.dNeighborSelf[i] = -1;
        } 
        if (e.prop.dNeighborSelf[i] == -(L - 1)) {
          e.prop.dNeighborSelf[i] = 1;
        }
      }

      unsigned vs = v1 < v2 ? v1 : v2;
      unsigned vl = v1 < v2 ? v2 : v1;

      if (v1 / L == v2 / L) {
        if (vl % L == L - 1 && vs % L == 0) {
          e.prop.spin = &(si.s[vl]);
        } else {
          e.prop.spin = &(si.s[vs]);
        }
      } else {
        if (vl / L == L - 1 && vs / L == 0) {
          e.prop.spin = &(si.s[pow(L, 2) + vl]);
        } else {
          e.prop.spin = &(si.s[pow(L, 2) + vs]);
        }
      }
    }
  }

  // add the next-nearest-neighbor bonds
  for (graph_x::vertex &v : Gxy.vertices) {
    int i = v.ind / L;
    int j = v.ind % L;

    graph_x::halfedge e1(v, Gxy.vertices[((i - 1) % L) * L + (j + 1) % L]);
    graph_x::halfedge e2(v, Gxy.vertices[((i - 1) % L) * L + (j - 1) % L]);
    graph_x::halfedge e3(v, Gxy.vertices[((i + 1) % L) * L + (j + 1) % L]);
    graph_x::halfedge e4(v, Gxy.vertices[((i + 1) % L) * L + (j - 1) % L]);

    e1.prop.distance = 1;
    e2.prop.distance = 1;
    e3.prop.distance = 1;
    e4.prop.distance = 1;

    e1.prop.dNeighborSelf[0] =  1;
    e1.prop.dNeighborSelf[1] = -1;

    e2.prop.dNeighborSelf[0] = -1;
    e2.prop.dNeighborSelf[1] = -1;

    e3.prop.dNeighborSelf[0] =  1;
    e3.prop.dNeighborSelf[1] =  1;

    e4.prop.dNeighborSelf[0] = -1;
    e4.prop.dNeighborSelf[1] =  1;

    v.edges.push_back(e1);
    v.edges.push_back(e2);
    v.edges.push_back(e3);
    v.edges.push_back(e4);
  }

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

  // now that the xy spins have been initialized, go back to the ising graph
  // and add a reference to the ones on each bond
  for (unsigned i = 0; i < 2; i++) {
    unsigned sb = i * pow(L, 2);

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

      unsigned xyv1 = j;
      unsigned xyv2;
      if (i == 0) {
        xyv2 = L * (xyv1 / L) + ((xyv1 % L) + 1) % L;
      } else {
        xyv2 = L * (((xyv1 / L) + 1) % L) + (xyv1 % L);
      }

      si.G.vertices[vc].prop.back = &(sxy.s[xyv1]);
      si.G.vertices[vc].prop.front = &(sxy.s[xyv2]);
    }
  }

  // start in the ground state of the Ising model
  for (unsigned i = 0; i < pow(L, D); i++) {
    si.s[pow(L, D) + i].x = true;
  }

  cos_xy = K1 * 2 * pow(L, 2) + K1 * K2 * 4 * pow(L, 2);
  sin_xy.fill(0);

  i_measurement mi(si, α, K1, nh, w);
  x_measurement mxy(sxy, K1, K2, α, &(si.s0));

  quantity_pair stiffness(nh, T);

  unsigned nn = 0;

  while (true) {
    si.run_wolff(1, ising::wolff::gen_ising<graph_i>, mi, rng);
    sxy.run_wolff(1, xy::wolff::generate_rotation_uniform<2, graph_x>, mxy, rng);
    if (nn > w) {
      stiffness.add(cos_xy, sin_xy * sin_xy);

      if (nn % N == 0) {
        double err = mi.A2.serr();
        double val2 = stiffness.avg();
        double err2 = stiffness.serr();
        if (err / mi.A2.avg() <= εmag && err2 / val2 < εstiff) {
          break;
        }
      }
    }
    nn++;
  }

  std::ifstream checkfile("out.dat");
  bool already_exists = checkfile.good();
  if (already_exists) {
    checkfile.close();
  }

  std::ofstream outfile;
  outfile.open("out.dat", std::ios::app);
  if (!already_exists) {
    outfile << "N L T J K1 K2 \\[Alpha] M \\[Rho] \\[Sigma]\n";
  }
  outfile << N << " " << L << " " << T << " " << J << " " << K1 << " " << K2 << " " << α << " " << mi.A2.avg() / pow(si.nv, 2) << " " << mi.A2.serr() / pow(si.nv, 2) << " " << stiffness.avg() / sxy.ne << " " << stiffness.serr() / sxy.ne << "\n";
  outfile.close();
}