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

#include "randutils/randutils.hpp"

#define WOLFF_USE_FINITE_STATES
#include "wolff/lib/wolff_models/ising.hpp"

using namespace wolff;

class track : public measurement<ising_t, ising_t, graph<>> {
public:
  int e;
  int m;
  track(const wolff::system<ising_t, ising_t, graph<>>& s) {
    e = -s.ne;
    m = s.nv;
  }

  void ghost_bond_visited(const system<ising_t, ising_t, graph<>>&, const typename graph<>::vertex&,
                          const ising_t& s_old, const ising_t& s_new, double dE) override {
    m += s_new - s_old;
  }

  void plain_bond_visited(const wolff::system<ising_t, ising_t, graph<>>& s,
                          const typename graph<>::halfedge& ed, const ising_t& si_new,
                          double dE) override {
    e -= 2 * (si_new * s.s[ed.neighbor.ind]);
  }
};

class sample : public measurement<ising_t, ising_t, graph<>> {
private:
  int e;
  int m;
  int64_t e_tot;
  int64_t m_tot;
  unsigned N;

public:
  sample(int eold, int mold) {
    e = eold;
    m = mold;
    e_tot = 0;
    m_tot = 0;
    N = 0;
  }

  void ghost_bond_visited(const system<ising_t, ising_t, graph<>>&, const typename graph<>::vertex&,
                          const ising_t& s_old, const ising_t& s_new, double dE) override {
    m += s_new - s_old;
  }

  void plain_bond_visited(const wolff::system<ising_t, ising_t, graph<>>& s,
                          const typename graph<>::halfedge& ed, const ising_t& si_new,
                          double dE) override {
    e -= 2 * (si_new * s.s[ed.neighbor.ind]);
  }

  void post_cluster(unsigned n, unsigned,
                    const wolff::system<ising_t, ising_t, graph<>>&) override {
    e_tot += e;
    m_tot += m;
    N++;
  }

  double e_avg() const {
    return ((double)e_tot) / N;
  }

  double m_avg() const {
    return ((double)m_tot) / N;
  }
};

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

  // set defaults
  unsigned N = (unsigned)1e4;
  unsigned D = 2;
  unsigned L = 128;
  double T = 2 / log(1 + sqrt(2));
  double H = 0;
  double w = 1.0;

  int opt;

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

  unsigned N_wait = w * pow(L, D);

  // define the spin-spin coupling
  std::function<double(const ising_t&, const ising_t&)> Z =
      [](const ising_t& s1, const ising_t& s2) -> double { return (double)(s1 * s2); };

  // define the spin-field coupling
  std::function<double(const ising_t&)> B = [=](const ising_t& s) -> double { return H * s; };

  // initialize the lattice
  graph<> G(D, L);

  // initialize the system
  wolff::system<ising_t, ising_t, graph<>> S(G, T, Z, B);

  // initialize the random number generator
  randutils::auto_seed_128 seeds;
  std::mt19937 rng{seeds};

  // define function that generates self-inverse rotations
  std::function<ising_t(std::mt19937&, const wolff::system<ising_t, ising_t, graph<>>&,
                        const graph<>::vertex&)>
      gen_r = gen_ising<graph<>>;

  // initailze the measurement object
  track AA(S);
  S.run_wolff(N_wait, gen_r, AA, rng);
  sample A(AA.e, AA.m);
  S.run_wolff(N, gen_r, A, rng);

  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 D L T H E M\n";
  }
  outfile << N << " " << D << " " << L << " " << T << " " << H << " " << A.e_avg() / S.nv << " " << A.m_avg() / S.nv << "\n";
  outfile.close();

  // exit
  return 0;
}