#include <vector>
#include <stack>

#include "randutils/randutils.hpp"
#include "pcg-cpp/include/pcg_random.hpp"

#include "blossom5-v2.05.src/PerfectMatching.h"

using Rng = randutils::random_generator<pcg32>;
using Real = long double;

class AztecDiamond {
public:
  using Coordinate = std::array<int, 2>;

  typedef struct Vertex {
    unsigned index;
    Coordinate coordinate;
  } Vertex;

  typedef struct Edge {
    Vertex* tail;
    Vertex* head;
    Real weight;
    std::stack<Real> weights;
    Real probability = 0;
  } Edge;

private:
  std::tuple<Edge&, Edge&, Edge&, Edge&> face(unsigned i, unsigned j) {
    unsigned x0 = n - i;
    unsigned x = x0 + 2 * (j % i);
    unsigned y = x0 + 2 * (j / i);

    Edge& e1 = edges[2 * n * y + x];
    Edge& e2 = edges[2 * n * y + x + 1];
    Edge& e3 = edges[2 * n * (y + 1) + x];
    Edge& e4 = edges[2 * n * (y + 1) + x + 1];

    return {e1, e2, e3, e4};
  }

public:
  unsigned n;
  std::vector<Vertex> vertices;
  std::vector<Edge> edges;

  AztecDiamond(int n) : n(n), vertices(2 * n * (n + 1)), edges(pow(2 * n, 2)) {
    unsigned M = vertices.size() / 2;
    for (int i = 0; i < M; i++) {
      vertices[i].index = i;
      vertices[M + i].index = M + i;
      vertices[i].coordinate = {2 * (i % (n + 1)), 2 * (i / (n + 1)) + 1};
      vertices[M + i].coordinate = {2 * (i % n) + 1, 2 * (i / n)};
    }
    for (unsigned i = 0; i < edges.size(); i++) {
      edges[i].tail = &vertices[(1 + (i % (2 * n))) / 2 + (n + 1) * ((i / 4) / n)];
      edges[i].head = &vertices[M + (i % (2 * n)) / 2 + n * (((i + 2 * n) / 4) / n)];
    }
  }

  void setWeights(Rng& r) {
    for (Edge& e : edges) {
      e.weight = r.variate<Real, std::exponential_distribution>(1);
    }
  }

  void computeWeights(Real T) {
    for (Edge& e : edges) {
      e.weights.push(exp(-e.weight / T));
    }

    for (unsigned i = n; i > 0; i--) {
#pragma omp parallel for
      for (unsigned j = 0; j < pow(i, 2); j++) {
        auto [e1, e2, e3, e4] = face(i, j);

        Real w = e1.weights.top();
        Real x = e2.weights.top();
        Real y = e3.weights.top();
        Real z = e4.weights.top();

        Real cellFactor = w * z + x * y;

        e1.weights.push(z / cellFactor);
        e2.weights.push(y / cellFactor);
        e3.weights.push(x / cellFactor);
        e4.weights.push(w / cellFactor);
      }
    }

    // This process computes one extra weight per edge.
    for (Edge& e : edges) {
      e.weights.pop();
    }
  }

  Real computeProbabilities() { // destroys *all* weights
    for (Edge& e : edges) {
      e.probability = 0;
    }
    Real logPartitionFunction = 0;

    for (unsigned i = 1; i <= n; i++) {
#pragma omp parallel for reduction(+:logPartitionFunction)
      for (unsigned j = 0; j < pow(i, 2); j++) {
        auto [e1, e2, e3, e4] = face(i, j);

        Real p = e1.probability;
        Real q = e2.probability;
        Real r = e3.probability;
        Real s = e4.probability;

        Real w = e1.weights.top();
        Real x = e2.weights.top();
        Real y = e3.weights.top();
        Real z = e4.weights.top();

        Real cellFactor = w * z + x * y;
        Real deficit = 1 - p - q - r - s;

        e1.probability = s + deficit * w * z / cellFactor;
        e2.probability = r + deficit * x * y / cellFactor;
        e3.probability = q + deficit * x * y / cellFactor;
        e4.probability = p + deficit * w * z / cellFactor;

        e1.weights.pop();
        e2.weights.pop();
        e3.weights.pop();
        e4.weights.pop();

        logPartitionFunction += log(cellFactor);
      }
    }

    return logPartitionFunction;
  }
};