#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;
  unsigned nc;

public:
  std::vector<double> m;

  Tree(unsigned n, unsigned n_moments) : p(n, -1), o(n, 1), m(n_moments, 1) {
    nc = 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) {
    for (unsigned k = 1; k <= m.size(); k++) {
      m[k - 1] -= (pow(o[i], k) + pow(o[j], k)) / nc;
    }

    nc--;

    if (o[i] < o[j]) {
      p[i] = j;
      o[j] += o[i];
      for (unsigned k = 1; k <= m.size(); k++) {
        m[k - 1] = ((nc + 1) * m[k - 1] + pow(o[j], k)) / nc;
      }
    } else {
      p[j] = i;
      o[i] += o[j];
      for (unsigned k = 1; k <= m.size(); k++) {
        m[k - 1] = ((nc + 1) * m[k - 1] + pow(o[i], k)) / nc;
      }
    }
  }

  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, unsigned n, const std::vector<std::vector<double>>& data,
                              std::string model_string) {
  std::string filename = model_string + id + ".dat";
  std::ifstream file(filename);

  unsigned n_old = 0;
  std::vector<std::vector<double>> data_old(data.size());

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

  if (file.is_open()) {
    file >> n_old;

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

    file.close();
  }

  std::ofstream file_out(filename);

  file_out << n_old + n << "\n";

  for (unsigned i = 0; i < data.size(); i++) {
    for (unsigned j = 0; j < data[i].size(); j++) {
      file_out << std::fixed << (data_old[i][j] * n_old + data[i][j]) / (n_old + n) << " ";
    }
    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};
  }

  unsigned n_moments = 5;

  std::vector<std::vector<double>> sn(n_moments);

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

  for (unsigned trial = 0; trial < N; trial++) {
    Tree t(n_vertices, n_moments);
    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 < n_moments; i++) {
        sn[i][p] += t.m[i];
      }
    }
  }

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

  return 0;
}