#include "hadamard_mcmc.hpp"
#include "quantity.hpp"
#include "matrices.hpp"

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

std::string getFilename(unsigned n, double β, double θ₀) {
  return "kauzmann_" + std::to_string(n) + "_" + std::to_string(β) + "_" + std::to_string(θ₀) + ".dat";
}

class MeasureEnergy : public Measurement {
public:
  Quantity Eavg;

  MeasureEnergy(unsigned lag) : Eavg(lag) {}

  void after_sweep(double, double E, const Orthogonal& M) override {
    Eavg.add(E);
  }
};

int main(int argc, char* argv[]) {
  // model parameters
  unsigned n = 2; // matrix size over four

  // temperatures to simulate
  double β₀ = 1; // starting temperature
  double Δβ = 0.5; // step size
  double β₁ = 50; // ending temperature

  // simulation settings
  unsigned m = 1e4; // number of tuning sweeps at each temperature
  unsigned N = 1e3; // number of autocorrelation times to simulate
  double ε = 0.1; // uncertainty on autocorrelation time
  double θ₀ = 0.05; // standard deviation of step size

  // measurement settings
  unsigned l = 1e3; // lag in autocorrelation function

  bool loadInitialFromFile = false;
  std::string inFilename;

  bool loadDataFromFile = false;

  int opt;

  while ((opt = getopt(argc, argv, "n:b:d:c:m:N:e:l:t:i:L")) != -1) {
    switch (opt) {
    case 'n':
      n = atoi(optarg);
      break;
    case 'b':
      β₀ = atof(optarg);
      break;
    case 'd':
      Δβ = atof(optarg);
      break;
    case 'c':
      β₁ = atof(optarg);
      break;
    case 'm':
      m = (unsigned)atof(optarg);
      break;
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    case 'e':
      ε = atof(optarg);
      break;
    case 'l':
      l = (unsigned)atof(optarg);
      break;
    case 't':
      θ₀ = atof(optarg);
      break;
    case 'i':
      loadInitialFromFile = true;
      inFilename = optarg;
      break;
    case 'L':
      loadDataFromFile = true;
      break;
    default:
      exit(1);
    }
  }

  double β = β₀;

  MeasureEnergy energyMeasurement(l);
  MCMC simulation(n, β₀, energyMeasurement, θ₀);

  if (loadInitialFromFile) {
    std::ifstream inFile(inFilename);
    inFile.ignore(std::numeric_limits<std::streamsize>::max(),'\n');
    inFile.ignore(std::numeric_limits<std::streamsize>::max(),'\n');
    for (unsigned i = 0; i < n; i++) {
      for (unsigned j = 0; j < n; j++) {
        inFile >> simulation.M(i, j);
      }
    }

    simulation.E = simulation.M.energy();

    std::cout << "Kauzmann: Imported matrix from file, energy " << simulation.E << "." << std::endl;
  }

  if (loadDataFromFile) {
    energyMeasurement.Eavg.read(getFilename(n, β, θ₀));
    std::cout << "Kauzmann: Imported data from file, number of steps " << energyMeasurement.Eavg.num_added() << "." << std::endl;
  }

  while (β <= β₁) {
    std::cout << "Kauzmann: Starting simulation of " << β << "." << std::endl;
    double rat = 0;
    unsigned num = 0;
    simulation.run(m, true);
    simulation.run(l);
    while (true) {
      Quantity EavgBak = energyMeasurement.Eavg;
      Orthogonal Mbak = simulation.M;
      rat += simulation.run(m);
      num++;
      std::array<double, 2> τ = energyMeasurement.Eavg.τ();
      std::cout << "Kauzmann: After " << energyMeasurement.Eavg.num_added() << " steps the autocorrelation time is " << τ[0] << " ± " << τ[1] << "." << std::endl;
      if (τ[1] / τ[0] < ε && τ[0] * N <= energyMeasurement.Eavg.num_added()) {
        break;
      }
    }
    std::cout << "Kauzmann: Finished simulation of " << β << ", " << rat / num << " acceptance ratio, writing output file." << std::endl;
    std::string filename = getFilename(n, β, θ₀);
    energyMeasurement.Eavg.write(filename);
    std::ofstream outFile(filename, std::ios::app);
    for (unsigned i = 0; i < n; i++) {
      for (unsigned j = 0; j < n; j++) {
        outFile << simulation.M(i, j) << " ";
      }
    }
    outFile << std::endl;
    outFile.close();
    energyMeasurement.Eavg.reset();
    β += Δβ;
    simulation.β = β;
  }

  return 0;
}