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

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

using Rng = randutils::random_generator<pcg32>;

using Bool = short unsigned;
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:
  unsigned index;
  std::vector<HalfEdge> outgoingEdges;
  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 i₁, unsigned i₂) const {
    unsigned Δx² = 0;
    for (unsigned d = 0; d < D; d++) {
      int x₁ = vertices[i₁].coordinates[d];
      int x₂ = vertices[i₂].coordinates[d];
      unsigned Δx = abs(x₁ - x₂);
      if (Δx > L / 2) {
        Δx = L - Δx;
      }
      Δx² += Δx * Δx;
    }
    return Δx²;
  }

  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].index = i;
      vertices[i].coordinates.reserve(D);
      vertices[i].outgoingEdges.reserve(2 * D);
      for (unsigned d = 0; d < D; d++) {
        vertices[i].coordinates.push_back((i / uPow(L, d)) % L);

        unsigned j₁ = uPow(L, d + 1) * (i / uPow(L, d + 1)) + (i + uPow(L, d)) % uPow(L, d + 1);
        unsigned j₂ = uPow(L, d + 1) * (i / uPow(L, d + 1)) + (uPow(L, d + 1) + i - uPow(L, d)) % uPow(L, d + 1);

        unsigned e₁ = d * Nv() + i;
        unsigned e₂ = d * Nv() + j₂;

        HalfEdge he₁(e₁, vertices[j₁]);
        HalfEdge he₂(e₂, vertices[j₂]);

        vertices[i].outgoingEdges.push_back(he₁);
        vertices[i].outgoingEdges.push_back(he₂);
      }
    }

    for (const Vertex& v₁ : vertices) {
      for (const Vertex& v₂ : vertices) {
        unsigned Δx² = squaredDistance(v₁.index, v₂.index);
        if (Δx² > 0) {
          multiplicities[Δx² - 1]++;
        }
      }
    }
  }

  std::list<Cluster> findClusters(const std::vector<Bool>& activeEdges) const {
    std::list<Cluster> clusters;
    std::vector<unsigned> clusterIndicies(Nv());
    unsigned nClusters = 0;
    for (const Vertex& v₀ : vertices) {
      if (clusterIndicies[v₀.index] == 0) {
        nClusters++;
        clusters.push_back({});
        std::stack<std::reference_wrapper<const Vertex>> verticiesToProcess;
        verticiesToProcess.push(v₀);
        while (!verticiesToProcess.empty()) {
          const Vertex& v₁ = verticiesToProcess.top();
          verticiesToProcess.pop();
          if (clusterIndicies[v₁.index] == 0) {
            clusterIndicies[v₁.index] = nClusters;
            clusters.back().push_back(v₁);
            for (const HalfEdge& e : v₁.outgoingEdges) {
              if (activeEdges[e.index] && clusterIndicies[e.neighbor.index] == 0) {
                verticiesToProcess.push(e.neighbor);
              }
            }
          }
        }
      }
    }

    return clusters;
  }

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

    for (const Vertex& v : vertices) {
      for (const HalfEdge& e : v.outgoingEdges) {
        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; // Dimension of space
  unsigned L = 8; // Size of lattice
  Real p = 0.5;   // Probability of bond
  unsigned M = 1; // Moments of conductivity
  unsigned n = 1; // Number of experiments

  int opt;

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

  Graph G(D, L);

  for (unsigned trial = 0; trial < n; trial++) {
    std::vector<Bool> activeEdges(G.Ne());
    for (Bool& edge : activeEdges) {
      edge = r.uniform(0.0, 1.0) < p;
    }

    std::list<Cluster> clusters = G.findClusters(activeEdges);
    Matrix laplacian = G.laplacian(activeEdges);

    std::vector<std::vector<Real>> conductivities(M + 1);

    for (std::vector<Real>& dat : conductivities) {
      dat.resize((L * L * D) / 4);
    }

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

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

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

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

          for (unsigned m = 0; m <= M; m++) {
            conductivities[m][G.squaredDistance(clusterIndicies[i], clusterIndicies[j]) - 1] += pow(1 / ΔV, m);
          }
        }
      }
    }

    for (unsigned m = 0; m <= M; m++) {
      std::ofstream file("perc_" + std::to_string(D) + "_" + std::to_string(L) + "_" + std::to_string(p) + "_" + std::to_string(m) + ".dat", std::ios_base::app);
      for (unsigned i = 0; i < conductivities[m].size(); i++) {
        if (G.multiplicities[i] != 0) {
          file << conductivities[m][i] / G.multiplicities[i] << " ";
        } else {
          file << 0 << " ";
        }
      }
      file << std::endl;
    }
  }

  return 0;
}