#include "hadamard_pt.hpp"

#include <fstream>
#include <iostream>

class MeasureEnergy : public Measurement {
public:
  unsigned N;
  double totalE;

  MeasureEnergy() {
    N = 0;
    totalE = 0;
  }

  void after_sweep(double, double E, const Orthogonal&) override {
    totalE += E;
    N++;
  }

  double energy() const { return totalE / N; }
};

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

  MeasureTransitionRates(unsigned n) : nAccepted(n - 1) {
    total_steps = 0;
    for (unsigned i = 0; i < n - 1; i++) {
      nAccepted[i].resize(i + 1);
    }
  }

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

  void after_sweep(const std::vector<MCMC>&) override { total_steps++; }
};

int main(int argc, char* argv[]) {
  unsigned n_tuning = 1e2;
  double β0 = 1.00;
  double βf = 40.00;
  unsigned num = 40;
  unsigned k = 2;
  double ε = 0.01;

  unsigned M = 10;
  unsigned N = 1e4;

  int opt;

  while ((opt = getopt(argc, argv, "k:b:c:n:t:N:M:e:")) != -1) {
    switch (opt) {
    case 'k':
      k = atoi(optarg);
      break;
    case 'b':
      β0 = atof(optarg);
      break;
    case 'c':
      βf = atof(optarg);
      break;
    case 'e':
      ε = atof(optarg);
      break;
    case 'n':
      num = 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;
    default:
      exit(1);
    }
  }

  unsigned n = pow(2, k);

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

  PT p(β0, βf, num, n, B, As);

  for (MCMC& sim : p.Ms) {
    sim.M = walsh(k);
  }

  std::cout << "Beginning " << n_tuning << " tuning steps for " << n << ".\n";
  p.tune(n_tuning, ε);
  std::cout << "Finished tuning, beginning " << N << " tempering updates of " << M
            << " sweeps each.\n";
  p.run(N, M);
  std::cout << "Finished " << n << ".\n";

  std::string filename = "probs_" + std::to_string(n) + "_" + std::to_string(β0) + "_" +
                         std::to_string(βf) + "_" + std::to_string(num) + ".dat";
  std::ifstream file(filename);

  unsigned N_old = 0;
  std::vector<std::vector<unsigned long>> data_old(B.nAccepted.size());

  for (unsigned i = 0; i < B.nAccepted.size(); i++) {
    data_old[i].resize(B.nAccepted[i].size());
  }

  if (file.is_open()) {
    file >> N_old;

    for (unsigned i = 0; i < B.nAccepted.size(); i++) {
      for (unsigned j = 0; j < B.nAccepted[i].size(); j++) {
        double num;
        file >> num;
        data_old[i][j] = num;
      }
    }

    file.close();
  }

  std::ofstream file_out(filename);

  file_out << N_old + B.total_steps << "\n";

  for (unsigned i = 0; i < B.nAccepted.size(); i++) {
    for (unsigned j = 0; j < B.nAccepted[i].size(); j++) {
      file_out << std::fixed << data_old[i][j] + B.nAccepted[i][j] << " ";
    }
    file_out << "\n";
  }

  file_out.close();

  std::string efilename = "energies_" + std::to_string(n) + "_" + std::to_string(β0) + "_" +
                          std::to_string(βf) + "_" + std::to_string(num) + ".dat";
  std::ifstream efile(efilename);

  unsigned Ne_old = 0;
  std::vector<double> edata_old(p.As.size());

  if (efile.is_open()) {
    efile >> N_old;

    for (unsigned i = 0; i < p.As.size(); i++) {
      double num;
      efile >> num;
      edata_old[i] = num;
    }

    efile.close();
  }

  std::ofstream efile_out(efilename);

  efile_out << Ne_old + ((MeasureEnergy*)As[0])->N << "\n";

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

  efile_out.close();

  return 0;
}