#include "hadamard_pt.hpp"
#include "matrices.hpp"

#include <execution>
#include <chrono>
#include <fstream>
#include <iostream>

class MeasureEnergy : public Measurement {
public:
  unsigned N;
  double totalE;
  double totalE2;
  unsigned n;
  std::vector<unsigned> ρ_dist;

  MeasureEnergy(unsigned n_bins = 1e4) : ρ_dist(n_bins + 1, 0) {
    n = n_bins;
    N = 0;
    totalE = 0;
    totalE2 = 0;
  }

  void after_sweep(double, double E, const Orthogonal& M) override {
    N++;
    totalE += E;
    totalE2 += pow(E, 2);
    double max = sqrt(M.size());
    for (unsigned i = 0; i < M.size(); i++) {
      for (unsigned j = 0; j < M.size(); j++) {
        ρ_dist[n * (M(i, j) + max) / (2 * max)]++;
      }
    }
  }

  double energy() const { return totalE / N; }
  double specific_heat() const { return totalE2 / N - pow(totalE / N, 2); }
};

class MeasureTransitionRates : public ParallelMeasurement {
public:
  std::vector<unsigned> nAccepted;
  std::vector<unsigned> total_steps;

  MeasureTransitionRates(unsigned n) : nAccepted(n - 1, 0), total_steps(n - 1, 0) {}

  void after_step(bool accepted, unsigned i, double, double, const MCMC&, const MCMC&) override {
    total_steps[i]++;
    if (accepted)
      nAccepted[i]++;
  }
};

int main(int argc, char* argv[]) {
  unsigned n_tuning = 1e2;
  double β₀ = 0.1;
  double β₁ = 10;
  unsigned N = 16;
  unsigned k = 2;
  double ε = 0.01;
  double ε2 = 0.01;

  unsigned M = 10;
  unsigned m = 1e4;

  bool loadBetasFromFile = false;
  std::string betaFilename;

  int opt;

  while ((opt = getopt(argc, argv, "k:b:c:n:t:N:M:e:m:f:F:")) != -1) {
    switch (opt) {
    case 'k':
      k = atoi(optarg);
      if (k == 0 || k > 8) {
        std::cout << "The size k must be an integer from 1 to 8!" << std::endl;
        exit(1);
      }
      break;
    case 'b':
      β₀ = atof(optarg);
      break;
    case 'c':
      β₁ = atof(optarg);
      break;
    case 'e':
      ε = atof(optarg);
      break;
    case 'f':
      ε2 = atof(optarg);
      break;
    case 'n':
      m = (unsigned)atof(optarg);
      break;
    case 't':
      n_tuning = (unsigned)atof(optarg);
      break;
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    case 'M':
      M = (unsigned)atof(optarg);
      break;
    case 'F':
      loadBetasFromFile = true;
      betaFilename = optarg;
      break;
    default:
      exit(1);
    }
  }

  unsigned n = 4 * k;

  std::vector<Measurement*> As(N);
  for (Measurement*& A : As) {
    A = new MeasureEnergy();
  }
  MeasureTransitionRates B(N);

  PT p(β₀, β₁, N, n, B, As);

  for (MCMC& sim : p.Ms) {
    sim.M = hadamards[k - 1];
    sim.E = sim.M.energy();
  }

  std::vector<double> f;

  std::cout << "Beginning simulation of " << n << ".\n";
  if (loadBetasFromFile) {
    std::cout << "Loading βs from file.\n";
    std::ifstream betaFile(betaFilename);
    for (unsigned i = 0; i < N; i++) {
      double βtmp;
      betaFile >> βtmp;
      p.Ms[i].β = βtmp;
    }
    std::cout << "βs: ";
    for (const MCMC& M : p.Ms) {
      std::cout << M.β << " ";
    }
    std::cout << std::endl;
    std::cout << "Beginning " << n_tuning << " tuning tempering updates.\n";
    Rng rng;
    for (unsigned i = 0; i < n_tuning; i++) {
      std::for_each(std::execution::par_unseq, p.Ms.begin(), p.Ms.end(),
                    [M, ε](MCMC& MM) { MM.tune(M, ε); });
      for (unsigned j = 0; j < p.Ms.size(); j++) {
        unsigned k = rng.uniform((unsigned)0, (unsigned)(p.Ms.size() - 2));
        p.step(k, k + 1, true);
      }
    }
  } else {
    std::cout << "Beginning " << n_tuning << " tuning tempering updates of " << M
              << " sweeps each.\n";
    f = p.tune(n_tuning, M, ε, ε2);
    std::cout << "βs: ";
    for (const MCMC& M : p.Ms) {
      std::cout << M.β << " ";
    }
    std::cout << std::endl;
  }
  std::cout << "Running " << m << " PT swaps of " << M << " sweeps each.\n";
  p.run(m, M); 
  std::cout << "Finished " << n << ".\n";

  auto tag = std::chrono::high_resolution_clock::now();

  std::string filename = "hmm_" + std::to_string(n) + "_" + std::to_string(β₀) + "_" +
                         std::to_string(β₁) + "_" + std::to_string(N) + "_" +
                         std::to_string(tag.time_since_epoch().count()) + ".dat";

  std::ofstream file(filename);

  for (const MCMC& M : p.Ms) {
    file << M.β << " ";
  }
  file << std::endl;

  for (double ff : f) {
    file << ff << " ";
  }
  file << std::endl;

  for (unsigned i = 0; i < B.nAccepted.size(); i++) {
    file << std::fixed << B.nAccepted[i] / (double)B.total_steps[i] << " ";
  }
  file << std::endl;

  for (unsigned i = 0; i < As.size(); i++) {
    file << std::fixed << ((MeasureEnergy*)As[i])->energy() << " ";
  }

  file << std::endl;

  for (unsigned i = 0; i < As.size(); i++) {
    file << std::fixed << ((MeasureEnergy*)As[i])->specific_heat() << " ";
  }

  file << std::endl;

  for (unsigned i = 0; i < As.size(); i++) {
    for (unsigned j = 0; j < ((MeasureEnergy*)As[i])->ρ_dist.size(); j++) {
      file << std::fixed << ((MeasureEnergy*)As[i])->ρ_dist[j] << " ";
    }
    file << std::endl;
  }

  file.close();

  return 0;
}