#include <algorithm>
#include <cinttypes>
#include <fstream>
#include <random>
#include <string>
#include <vector>

#include "randutils/randutils.hpp"

class Tree {
private:
  std::vector<signed> p;
  std::vector<unsigned> o;

public:
  unsigned largest;
  std::vector<unsigned> c;

  Tree(unsigned n) : p(n, -1), o(n, 1), c(n, 0) {
    largest = 1;
    c[0] = n;
  }

  unsigned findroot(unsigned i) {
    if (p[i] < 0)
      return i;
    else
      return p[i] = this->findroot(p[i]);
  }

  void join(unsigned i, unsigned j) {
    c[o[i] - 1]--;
    c[o[j] - 1]--;

    if (o[i] < o[j]) {
      p[i] = j;
      o[j] += o[i];
      c[o[j] - 1]++;
      if (o[j] > largest) {
        largest = o[j];
      }
    } else {
      p[j] = i;
      o[i] += o[j];
      c[o[i] - 1]++;
      if (o[i] > largest) {
        largest = o[i];
      }
    }
  }

  void add_bond(const std::array<unsigned, 2>& b) {
    unsigned i = this->findroot(b[0]);
    unsigned j = this->findroot(b[1]);

    if (i != j) {
      this->join(i, j);
    }
  }
};

void update_distribution_file(std::string id, const std::vector<std::vector<uint64_t>>& data,
                              std::string model_string) {
  std::string filename = model_string + id + ".dat";
  std::ifstream file(filename);

  std::vector<std::vector<uint64_t>> data_old(data.size());

  for (unsigned i = 0; i < data.size(); i++) {
    data_old[i].resize(data[0].size());
  }

  if (file.is_open()) {
    for (unsigned i = 0; i < data.size(); i++) {
      for (unsigned j = 0; j < data[0].size(); j++) {
        uint64_t num;
        file >> num;
        data_old[i][j] = num;
      }
    }

    file.close();
  }

  std::ofstream file_out(filename);

  for (unsigned i = 0; i < data.size(); i++) {
    for (unsigned j = 0; j < data[0].size(); j++) {
      file_out << std::fixed << data_old[i][j] + data[i][j] << " ";
    }
    file_out << "\n";
  }

  file_out.close();
}

int main(int argc, char* argv[]) {
  unsigned L = 16;
  unsigned N = 100;

  int opt;

  while ((opt = getopt(argc, argv, "L:N:")) != -1) {
    switch (opt) {
    case 'L':
      L = (unsigned)atof(optarg);
      break;
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    default:
      exit(1);
    }
  }

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

  unsigned n_vertices = pow(L, 2);
  unsigned n_bonds = 2 * n_vertices;

  std::vector<std::array<unsigned, 2>> bonds(n_bonds);
  std::vector<unsigned> bond_indices(n_bonds);

  for (unsigned i = 0; i < n_bonds; i++) {
    bond_indices[i] = i;
  }

  for (unsigned v_i = 0; v_i < n_vertices; v_i++) {
    bonds[2 * v_i + 0] = {v_i, L * (v_i / L) + (v_i + 1) % L};
    bonds[2 * v_i + 1] = {v_i, L * (((v_i / L) + 1) % L) + v_i % L};
  }

  std::vector<std::vector<uint64_t>> sn(n_bonds);

  for (std::vector<uint64_t>& sni : sn) {
    sni.resize(n_vertices);
  }

  for (unsigned trial = 0; trial < N; trial++) {
    Tree t(n_vertices);
    std::shuffle(bond_indices.begin(), bond_indices.end(), rng);

    for (unsigned p = 0; p < n_bonds; p++) {
      t.add_bond(bonds[bond_indices[p]]);

      for (unsigned i = 0; i < t.largest; i++) {
        sn[p][i] += t.c[i];
      }
    }
  }

  update_distribution_file("sn", sn, "percolation_" + std::to_string(L) + "_");

  return 0;
}