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

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(std::list<range> ranges, unsigned n, ParallelMeasurement& B, std::vector<Measurement*>& As)
      : B(B), As(As) {
    unsigned count = 0;
    for (range r : ranges) {
      for (unsigned i = 1; i <= r.N; i++) {
        double β = r.β0 + i * (r.β1 - r.β0) / r.N;
        Ms.push_back(MCMC(n, β, *As[count]));
        count++;
      }
    }
  }

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

  bool step(unsigned i, unsigned j) {
    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]);

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

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

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