#include <functional>
#include <fstream>
#include <stack>

#include "rbmp.hpp"

using Real = long double;

class AztecDiamond {

public:
  typedef struct Edge {
    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 xx = 2 * (j % i);
    unsigned yy = 2 * (j / i);

    Edge& e1 = edges[2 * n * (x0 + yy) + x0 + xx];
    Edge& e2 = edges[2 * n * (x0 + yy) + x0 + xx + 1];
    Edge& e3 = edges[2 * n * (x0 + yy + 1) + x0 + xx];
    Edge& e4 = edges[2 * n * (x0 + yy + 1) + x0 + xx + 1];

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

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

  AztecDiamond(unsigned n) : edges(pow(2 * n, 2)), n(n) {}

  void computeWeights() {
    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
    Real partitionFunction = 1;

    for (unsigned i = 1; i <= n; i++) {
#pragma omp parallel for
      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();

        partitionFunction *= cellFactor;
      }
    }

    return partitionFunction;
  }
};

int main(int argc, char* argv[]) {
  unsigned n = 100;
  unsigned m = 100;
  Real T = 1;

  int opt;

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

  std::string filename = "order_" + std::to_string(n) + "_" + std::to_string(T) + ".dat";

  Rng r;
  AztecDiamond a(n);

  std::vector<Real> avgProbabilities(a.edges.size());

  for (unsigned i = 0; i < m; i++) {
    for (AztecDiamond::Edge& e : a.edges) {
      e.weights.push(exp(- r.variate<Real, std::exponential_distribution>(1) / T));
      e.probability = 0;
    }
    a.computeWeights();
    a.computeProbabilities();

    for (unsigned j = 0; j < a.edges.size(); j++) {
      avgProbabilities[j] += a.edges[j].probability;
    }
  }

  Graph G(n, r);

  std::vector<Real> data_x(G.vertices.size());
  std::vector<Real> data_y(G.vertices.size());

  for (unsigned i = 0; i < G.edges.size(); i++) {
    const Graph::Edge& e = G.edges[i];
    const Graph::Vertex& vt = e.halfedges[0].getTail();
    const Graph::Vertex& vh =  e.halfedges[0].getHead();
    data_x[vt.index] += avgProbabilities[i] * (vt.coordinate[0] - vh.coordinate[0]);
    data_y[vt.index] += avgProbabilities[i] * (vt.coordinate[1] - vh.coordinate[1]);
    data_x[vh.index] += avgProbabilities[i] * (vt.coordinate[0] - vh.coordinate[0]);
    data_y[vh.index] += avgProbabilities[i] * (vt.coordinate[1] - vh.coordinate[1]);
  }

  std::ofstream output(filename);

  for (unsigned i = 0; i < G.vertices.size(); i++) {
    output << data_x[i] << " " << data_y[i] << std::endl;
  }

  output.close();

  return 0;
}