#include "hadamard_pt.hpp"

#include <fstream>
#include <iostream>

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

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

  void after_sweep(double, double E, const Orthogonal& M) override {
    N++;
    totalE += E;
    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; }
};

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;
  std::list<double> β0s;
  std::list<double> β1s;
  std::list<unsigned> Ns;
  unsigned k = 2;
  double ε = 0.01;

  unsigned M = 10;
  unsigned N = 1e4;
  unsigned m = 1e4;

  int opt;

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

  unsigned n = pow(2, k);

  std::list<range> rs;
  unsigned num = 0;

  if ((β0s.size() != β1s.size()) || (β0s.size() != Ns.size())) {
    std::cout << "You need the same number of ranges!\n";
    exit(0);
  } else {
    auto it0 = β0s.begin();
    auto it1 = β1s.begin();
    auto itN = Ns.begin();

    while (it0 != β0s.end()) {
      rs.push_back({*it0, *it1, *itN});
      num += *itN;
      it0++;
      it1++;
      itN++;
    }
  }

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

  PT p(rs, n, B, As);

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

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

  std::string rs_string = "";

  for (range r : rs) {
    rs_string +=
        "_" + std::to_string(r.β0) + "_" + std::to_string(r.β1) + "_" + std::to_string(r.N);
  }

  std::string filename = "probs_" + std::to_string(n) + rs_string + ".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) + rs_string + ".dat";
  std::ifstream efile(efilename);

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

  if (efile.is_open()) {
    efile >> Ne_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();

  std::string ρfilename = "rhos_" + std::to_string(n) + rs_string + ".dat";
  std::ifstream ρfile(ρfilename);

  std::vector<std::vector<unsigned>> ρdata_old(As.size());

  for (unsigned i = 0; i < As.size(); i++) {
    ρdata_old[i].resize(((MeasureEnergy*)As[0])->ρ_dist.size());
  }

  if (ρfile.is_open()) {
    for (unsigned i = 0; i < As.size(); i++) {
      for (unsigned j = 0; j < ((MeasureEnergy*)As[0])->ρ_dist.size(); j++) {
        unsigned num;
        ρfile >> num;
        ρdata_old[i][j] = num;
      }
    }

    ρfile.close();
  }

  std::ofstream ρfile_out(ρfilename);

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

  ρfile_out.close();

  return 0;
}