#include "space_wolff.hpp"

int main(int argc, char* argv[]) {
  const unsigned D = 2;

  unsigned L = 32;
  unsigned N = 1000;
  double T = 2.0 / log(1.0 + sqrt(2.0));
  double H = 1.0;
  double ε = 0.1;

  int opt;

  while ((opt = getopt(argc, argv, "N:L:T:H:e:")) != -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;
      default:
        exit(1);
    }
  }

  std::function<double(spin<signed, D, signed>, spin<signed, D, signed>)> Z =
    [] (spin<signed, D, signed> s1, spin<signed, D, signed> s2) -> double {
      bool one_one = false;
      bool many_ones = false;
      bool any_two = false;
      
      for (unsigned i = 0; i < D; i++) {
        unsigned diff = abs(s1.x(i) - s2.x(i));
        if (diff == 1 && !one_one) {
          one_one = true;
        } else if (diff == 1 && one_one) {
          many_ones = true;
          break;
        } else if (diff > 1) {
          any_two = true;
          break;
        }
      }

      if (!one_one && !any_two) {
        return -std::numeric_limits<double>::infinity();
      } else if (one_one && !many_ones && !any_two) {
        return s1.s * s2.s;
      } else {
        return 0;
      }
    };

  std::function<double(spin<signed, D, signed>)> B =
    [L, H] (spin<signed, D, signed> s) -> double {
      return H * s.s * smiley[s.x(1) * 16 / L][s.x(0) * 16 / L];
    };

  std::function<std::set<unsigned>(model<signed, D, signed>&, unsigned, spin<signed, D, signed>)> neighbors =
    [] (model<signed, D, signed>& m, unsigned i0, spin<signed, D, signed> s1) -> std::set<unsigned> {
      std::set<unsigned> nn;
      if (i0 < m.s.size()) {
        std::set<unsigned> nn0 = m.dict.neighbors(m.s[i0].x, 1);
        std::set<unsigned> nn1 = m.dict.neighbors(s1.x, 1);
        nn.insert(nn0.begin(), nn0.end());
        nn.insert(nn1.begin(), nn1.end());
        nn.insert(m.s.size());
      } else {
        for (unsigned i = 0; i < m.s.size(); i++) {
          nn.insert(i);
        }
      }
      return nn;
    };

  model<signed, D, signed> ising(L, Z, B, neighbors);

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

  std::uniform_int_distribution<unsigned> coin(0, 1);

  unsigned n = 0;
  unsigned up = 0;
  unsigned down = 0;
  for (unsigned i = 0; i < L; i++) {
    for (unsigned j = 0; j < L; j++) {
      if ((coin(rng) && up < pow(L, 2) / 2) || down >= pow(L, 2) / 2) {
        ising.s.push_back({{i, j}, 1});
        up++;
      } else {
        ising.s.push_back({{i, j}, -1});
        down++;
      }
      ising.dict.record<signed>({i, j}, n);
      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++;
    }
  }
  */

  ising.update_energy();

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

  std::vector<signed> output(pow(L, D));

  for (spin<signed, D, signed> s : ising.s) {
    spin<signed, D, signed> rs = ising.s0.inverse().act(s);
    output[L * rs.x(1) + rs.x(0)] = s.s;
  }

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

  for (unsigned i = 0; i < L; i++) {
    for (unsigned j = 0; j < L; j++) {
      unsigned out = output[L * i + j] == 1 ? 1 : 0;
      outfile << out << " ";
    }
    outfile << "\n";
  }
  outfile.close();

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

  std::cout << act[0] << " " << act[1] << "\n";

  return 0;
}