#pragma once
#include "hadamard_mcmc.hpp"
#include <list>
#include <fstream>
#include <iostream>

void swap(MCMC& s1, MCMC& s2) {
  std::swap(s1.M, s2.M);
  std::swap(s1.E, s2.E);
}

class ParallelMeasurement {
public:
  virtual void after_step(bool, unsigned, unsigned, double, double, const MCMC&, const MCMC&){};
  virtual void after_sweep(const std::vector<MCMC>&){};
};

typedef struct range {
  double β0;
  double β1;
  unsigned N;
} range;


class PT {
private:
  randutils::mt19937_rng rng;

public:
  std::vector<MCMC> Ms;
  ParallelMeasurement& B;
  std::vector<Measurement*>& As;

  PT(double β₀, double β₁, unsigned N, unsigned n, ParallelMeasurement& B, std::vector<Measurement*>& As)
      : B(B), As(As) {
    Ms.reserve(N);
    for (unsigned i = 0; i < N; i++) {
      double β = β₀ + i * (β₁ - β₀) / (N - 1);
      Ms.push_back(MCMC(n, β, *As[i]));
    }
  }

  bool step(unsigned i, unsigned j, bool dry = false) {
    double Δβ = Ms[i].β - Ms[j].β;
    double ΔE = Ms[i].E - Ms[j].E;

    bool accepted = Δβ * ΔE > 0 || exp(Δβ * ΔE) > rng.uniform((double)0.0, 1.0);

    if (accepted)
      swap(Ms[i], Ms[j]);

    if (!dry) {
      B.after_step(accepted, i, j, Δβ, ΔE, Ms[i], Ms[j]);
    }
    return accepted;
  }

  void sweep(bool dry = false) {

    for (unsigned i = 0; i < Ms.size() - 1; i++) {
      this->step(i, i + 1, dry);
    }
  }

  void tune(unsigned n, unsigned m, double ε) {
    std::vector<color> colors(Ms.size(), none);
    std::vector<unsigned> nu(Ms.size(), 0);
    std::vector<unsigned> nd(Ms.size(), 0);

    colors.front() = down;
    colors.back() = up;

    for (unsigned i = 0; i < n; i++) {
#pragma omp parallel for
      for (unsigned j = 0; j < Ms.size(); j++) {
        Ms[j].tune(m, ε);
      }

      for (unsigned j = 0; j < Ms.size() - 1; j++) {
        if (this->step(j, j + 1, true)) {
          std::swap(colors[j], colors[j + 1]);
          colors.front() = down;
          colors.back() = up;

          if (colors[j] == up) {
            nu[j]++;
          } else if (colors[j] == down) {
            nd[j]++;
          }
          if (colors[j + 1] == up) {
            nu[j + 1]++;
          } else if (colors[j + 1] == down) {
            nd[j + 1]++;
          }
        }
      }
    }

    std::ofstream outfile("test.dat", std::ios::app);

    std::vector<double> f(Ms.size());

    double δ = 0.5;
    std::vector<unsigned> f_keep;
    double f_last = 0;
    for (unsigned i = 0; i < Ms.size(); i++) {
      f[i] = nu[i] / (double)(nu[i] + nd[i]);
      if (fabs(f[i] - (1 - (Ms.size() - i - 1) / (Ms.size() - 1.0))) < δ && f[i] >= f_last) {
        f_keep.push_back(i);
        f_last = f[i];
      }
    }

    for (unsigned i = 0; i < f_keep.size() - 1; i++) {
      for (unsigned j = f_keep[i]; j < f_keep[i + 1]; j++) {
        f[j] = f[f_keep[i]] + (f[f_keep[i + 1]] - f[f_keep[i]]) / (Ms[f_keep[i+1]].β - Ms[f_keep[i]].β) * (Ms[j].β - Ms[f_keep[i]].β);
      }
    }



    std::vector<double> η(Ms.size() - 1);
    double C = 0;

    for (unsigned i = 0; i < Ms.size() - 1; i++) {
      η[η.size() - i - 1] = sqrt((f[i + 1] - f[i]) / pow((1 / Ms[i].β - 1 / Ms[i + 1].β), 2));
    }

    std::vector<double> T0(Ms.size());
    std::vector<double> T1(Ms.size());

    for (unsigned i = 0; i < Ms.size(); i++) {
      T0[i] = 1 / Ms[Ms.size() - i - 1].β;
    }
    T1.front() = 1 / Ms.back().β;
    T1.back() = 1 / Ms.front().β;

    for (unsigned i = 0; i < η.size(); i++) {
      C += η[i] * (T0[i + 1] - T0[i]);
    }

    double x= 0;
    unsigned j = 1;

    for (unsigned i = 0; i < η.size(); i++) {
      double xnew = x + η[i] * (T0[i + 1] - T0[i]) / C;
      while (j < xnew * η.size()) {
        T1[j] = T0[i] + (j / (double)η.size() - x) / η[i] * C;
        j++;
      }
      x = xnew;
    }

    for (unsigned i = 0; i < T1.size(); i++) {
      Ms[Ms.size() - i - 1].β = 1 / T1[i];
    }

    for (double ff : f) {
      outfile << ff << " ";
    }
    outfile << "\n";
    outfile.close();

  }

  void run(unsigned n, unsigned m, bool dry = false) {
    for (unsigned i = 0; i < n; i++) {
#pragma omp parallel for
      for (unsigned j = 0; j < Ms.size(); j++) {
        Ms[j].run(m, dry);
      }
      this->sweep(dry);
      if (!dry) {
        B.after_sweep(this->Ms);
      }
    }
  }
};