#include "space_wolff.hpp"

std::function<double(spin<signed, 2, signed>)> B_sin(unsigned L, unsigned n, double H) {
  return [n, H, L] (spin<signed, 2, signed> s) -> double {
    return H * s.s * cos(2 * M_PI * n * s.x(0) / ((double)L));
  };
}

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

  unsigned L = 32;
  unsigned N = 1000;
  unsigned mod = 0;
  unsigned multi = 1e4;
  double mag = 0.5;
  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:m:M:n:")) != -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 'n':
        mag = atof(optarg);
        break;
      default:
        exit(1);
    }
  }

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

  std::function<double(spin<signed, D, signed>, spin<signed, D, signed>, spin<signed, D, signed>)> Z =
    [L, pZ] (spin<signed, D, signed> s1, spin<signed, D, signed> s2, spin<signed, D, signed> s1_new) -> double {
      bool old_one_one = false;
      bool old_many_ones = false;
      bool old_any_two = false;
      bool new_one_one = false;
      bool new_many_ones = false;
      bool new_any_two = false;

      vector<signed, D> old_diff = diff<signed, D>(L, s1.x, s2.x);
      vector<signed, D> new_diff = diff<signed, D>(L, s1_new.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;
        }
        if (new_diff(i) == 1 && !new_one_one) {
          new_one_one = true;
        } else if (new_diff(i) == 1 && new_one_one) {
          new_many_ones = true;
        } else if (new_diff(i) > 1) {
          new_any_two = true;
        }
      }

      bool were_on_someone = !old_one_one && !old_any_two;
      bool are_on_someone = !new_one_one && !new_any_two;
      bool were_nearest_neighbors = old_one_one && !old_many_ones && !old_any_two;
      bool are_nearest_neighbors = new_one_one && !new_many_ones && !new_any_two;

      if (were_on_someone) {
        return 0.0;
      } else if (are_on_someone) {
        return 1.0;
      } else if (were_nearest_neighbors && are_nearest_neighbors) {
        return 0.0;
      } else if (were_nearest_neighbors) {
        if (s1.s * s2.s == 1) {
          return pZ;
        } else {
          return 0.0;
        }
      } else if (are_nearest_neighbors) {
        if (s1_new.s * s2.s == -1) {
          return pZ;
        } else {
          return 0.0;
        }
      } else {
        return 0.0;
      }
    };

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

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

  if (mod > 0) {
    B = B_sin(L, mod, H);
  } else {
    B = B_face;
  }

  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) * mag) || down >= pow(L, 2) * mag) {
        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 << 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;
}