#include <stack>
#include <list>
#include <cstdlib>
#include <getopt.h>

#include "eigen/Eigen/Dense"
#include "randutils/randutils.hpp"
#include "pcg-cpp/include/pcg_random.hpp"

using Rng = randutils::random_generator<pcg32>;

using Real = double;
using Vector = Eigen::Matrix<Real, Eigen::Dynamic, 1>;
using Matrix = Eigen::Matrix<Real, Eigen::Dynamic, Eigen::Dynamic>;

class Vertex;

class HalfEdge {
public:
  unsigned index;
  Vertex& neighbor;
  HalfEdge(unsigned index, Vertex& v) : index(index), neighbor(v) {};
};

class Vertex {
public:
  std::vector<HalfEdge> bonds;
  unsigned index;
  std::vector<unsigned> coordinates;
};

class Cluster : public std::list<std::reference_wrapper<const Vertex>> {};

unsigned uPow(unsigned x, unsigned a) {
  unsigned result = 1;
  while (a) {
    if (a & 1) {
      result *= x;
    }
    a >>= 1;
    x *= x;
  }
  return result;
}

class Graph {
public:
  unsigned D;
  unsigned L;
  std::vector<Vertex> vertices;
  std::vector<unsigned> multiplicities;

  unsigned Nv() const {
    return vertices.size();
  }

  unsigned Ne() const {
    return D * Nv();
  }

  unsigned squaredDistance(unsigned vi, unsigned vj) const {
    unsigned sd = 0;
    for (unsigned i = 0; i < D; i++) {
      int x1 = vertices[vi].coordinates[i];
      int x2 = vertices[vj].coordinates[i];
      unsigned Δx = abs(x1 - x2);
      if (Δx > L / 2) {
        Δx = L - Δx;
      }
      sd += Δx * Δx;
    }
    return sd;
  }

  /* Initialize a square lattice */
  Graph(unsigned D, unsigned L) : D(D), L(L), vertices(uPow(L, D)), multiplicities((D * L * L) / 4) {
    for (unsigned i = 0; i < Nv(); i++) {
      vertices[i].coordinates.resize(D);
      vertices[i].index = i;
      vertices[i].bonds.reserve(2 * D);
      for (unsigned j = 0; j < D; j++) {
        vertices[i].coordinates[j] = (i / uPow(L, j)) % L;

        unsigned n1 = uPow(L, j + 1) * (i / uPow(L, j + 1)) + (i + uPow(L, j)) % uPow(L, j + 1);
        unsigned n2 = uPow(L, j + 1) * (i / uPow(L, j + 1)) + (uPow(L, j + 1) + i - uPow(L, j)) % uPow(L, j + 1);

        unsigned e1 = j * Nv() + i;
        unsigned e2 = j * Nv() + n2;

        HalfEdge he1(e1, vertices[n1]);
        HalfEdge he2(e2, vertices[n2]);

        vertices[i].bonds.push_back(he1);
        vertices[i].bonds.push_back(he2);
      }
    }

    for (const Vertex& v1 : vertices) {
      for (const Vertex& v2 : vertices) {
        unsigned dist = squaredDistance(v1.index, v2.index);
        if (dist > 0) {
          multiplicities[dist - 1]++;
        }
      }
    }
  }

  std::vector<Cluster> markClusters(const std::vector<short unsigned>& activeEdges) const {
    std::vector<Cluster> clusters;
    std::vector<unsigned> indicies(Nv());
    unsigned nClusters = 0;
    for (const Vertex& v : vertices) {
      if (indicies[v.index] == 0) {
        nClusters++;
        clusters.push_back({});
        std::stack<std::reference_wrapper<const Vertex>> inCluster;
        inCluster.push(v);
        while (!inCluster.empty()) {
          const Vertex& vn = inCluster.top();
          inCluster.pop();
          if (indicies[vn.index] == 0) {
            indicies[vn.index] = nClusters;
            clusters[nClusters - 1].push_back(vn);
          }
          for (const HalfEdge& e : vn.bonds) {
            if (activeEdges[e.index] && indicies[e.neighbor.index] == 0) {
              inCluster.push(e.neighbor);
            }
          }
        }
      }
    }

    return clusters;
  }

  Matrix laplacian(const std::vector<short unsigned>& activeEdges) const {
    Matrix L = Matrix::Zero(Nv(), Nv());

    for (const Vertex& v : vertices) {
      for (const HalfEdge& e : v.bonds) {
        if (activeEdges[e.index]) {
          L(v.index, v.index) += 1;
          L(v.index, e.neighbor.index) -= 1;
        }
      }
    }

    return L;
  }
};

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

  unsigned D = 2;
  unsigned L = 8;
  double p = 0.5;
  unsigned n = 10;

  int opt;

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

  Graph G(D, L);

  for (unsigned trials = 0; trials < n; trials++) {
    std::vector<short unsigned> activeEdges(G.Ne());
    for (short unsigned& edge : activeEdges) {
      edge = r.uniform(0.0, 1.0) < p;
    }
    std::vector<Cluster> clusters = G.markClusters(activeEdges);
    Matrix laplacian = G.laplacian(activeEdges);

    std::vector<Real> conductivities((L * L * D) / 4);
    std::vector<unsigned> probabilities((L * L * D) / 4);

    for (const Cluster& c : clusters) {
      std::vector<unsigned> inds;
      inds.reserve(c.size());
      for (const Vertex& v : c){
        inds.push_back(v.index);
      }

      Matrix subLaplacian = laplacian(inds, inds);
      subLaplacian(0,0)++; /* Set voltage of first node to zero */
      Matrix subLaplacianInv = subLaplacian.inverse();

      for (unsigned i = 0; i < c.size(); i++) {
        for (unsigned j = 0; j < i; j++) {
          Vector input = Vector::Zero(inds.size());
          input(i) = 1;
          input(j) = -1;

          Vector output = subLaplacianInv * input;
          double ΔV = std::abs(output(i) - output(j));

          conductivities[G.squaredDistance(inds[i], inds[j]) - 1] += 1 / ΔV;
          probabilities[G.squaredDistance(inds[i], inds[j]) - 1]++;
        }
      }
    }

    for (unsigned i = 0; i < conductivities.size(); i++) {
      if (G.multiplicities[i] != 0) {
        std::cout << conductivities[i] / G.multiplicities[i] << " ";
      } else {
        std::cout << 0 << " ";
      }
    }
    std::cout << std::endl;
    for (unsigned i = 0; i < probabilities.size(); i++) {
      if (G.multiplicities[i] != 0) {
        std::cout << ((Real)probabilities[i]) / G.multiplicities[i] << " ";
      } else {
        std::cout << 0 << " ";
      }
    }
    std::cout << std::endl;
  }

  return 0;
}