#pragma once
#include "pcg-cpp/include/pcg_random.hpp"
#include "randutils/randutils.hpp"
#include <vector>

using Rng = randutils::random_generator<pcg32>;

inline double Ei(double x) { return -fabs(x); }

class Orthogonal {
private:
  unsigned d;
  std::vector<double> m;

public:
  Orthogonal(unsigned size) : m(size * size) {
    d = size;
    for (unsigned i = 0; i < size; i++) {
      m[size * i + i] = sqrt(size);
    }
  }

  Orthogonal(const std::vector<double>& x) : m(x) { d = sqrt(x.size()); }

  unsigned size() const { return d; }

  double& operator()(unsigned i, unsigned j) { return m[d * i + j]; }

  const double& operator()(unsigned i, unsigned j) const { return m[d * i + j]; }

  double energy() const {
    double E = 0;

    for (unsigned i = 0; i < this->size(); i++) {
      for (unsigned j = 0; j < this->size(); j++) {
        E += Ei(this->operator()(i, j));
      }
    }

    return E;
  }

  void orthogonalize() {
    for (unsigned i = 0; i < d; i++) {
      for (unsigned j = 0; j < i; j++) {
        double ij = 0;
        for (unsigned k = 0; k < d; k++) {
          ij += (*this)(i, k) * (*this)(j, k);
        }

        ij /= d;

        for (unsigned k = 0; k < d; k++) {
          (*this)(i, k) -= ij * (*this)(j, k);
        }
      }

      double ii = 0;
      for (unsigned j = 0; j < d; j++) {
        ii += pow((*this)(i, j), 2);
      }

      ii = sqrt(ii / d);

      for (unsigned j = 0; j < d; j++) {
        (*this)(i, j) = (*this)(i, j) / ii;
      }
    }
  }
};

class Givens {
private:
  Orthogonal& m;

  bool transpose;
  unsigned axis_1;
  unsigned axis_2;
  double Δθ;

public:
  Givens(Orthogonal& m, bool t, unsigned a1, unsigned a2, double θ0, Rng& rng) : m(m) {
    transpose = t;
    axis_1 = a1;
    axis_2 = a2;
    Δθ = rng.uniform(-θ0, θ0);
  }

  Givens(Orthogonal& m, double θ0, Rng& rng) : m(m) {
    Δθ = rng.uniform(-θ0, θ0);
    unsigned axis1axis2 = rng.uniform((unsigned)0, m.size() * (m.size() - 1) - 1);

    axis_1 = axis1axis2 / (m.size() - 1);
    axis_2 = axis1axis2 % (m.size() - 1);
    transpose = axis_2 >= axis_1;
    if (transpose) {
      axis_2++;
    }
  }

  double tryRotation(std::vector<double>& rows) const {
    double ΔE = 0;
    double c = cos(Δθ);
    double s = sin(Δθ);

    for (unsigned i = 0; i < m.size(); i++) {
      double m1i, m2i, m1i_new, m2i_new;

      if (transpose) {
        m1i = m(i, axis_1);
        m2i = m(i, axis_2);
      } else {
        m1i = m(axis_1, i);
        m2i = m(axis_2, i);
      }

      ΔE -= Ei(m1i) + Ei(m2i);

      m1i_new = c * m1i + s * m2i;
      m2i_new = c * m2i - s * m1i;

      ΔE += Ei(m1i_new) + Ei(m2i_new);

      rows[i] = m1i_new;
      rows[m.size() + i] = m2i_new;
    }

    return ΔE;
  }

  void acceptRotation(std::vector<double>& rows) const {
    for (unsigned i = 0; i < m.size(); i++) {
      if (transpose) {
        m(i, axis_1) = rows[i];
        m(i, axis_2) = rows[m.size() + i];
      } else {
        m(axis_1, i) = rows[i];
        m(axis_2, i) = rows[m.size() + i];
      }
    }
  }
};

class Measurement {
public:
  virtual void after_step(bool, const Givens&, double, double, double, const Orthogonal&){};
  virtual void after_sweep(double, double, const Orthogonal&){};
};

typedef enum { none, up, down } color;

class MCMC {
private:
  std::vector<double> row_storage;
  Measurement& A;
  double θ0;

public:
  Rng rng;
  double β;
  double E;
  Orthogonal M;

  MCMC(unsigned n, double β0, Measurement& A) : A(A), M(n), β(β0), row_storage(2 * n) {
    θ0 = M_PI;
    E = M.energy();
  }

  bool step(Givens& g, bool dry = false) {
    double ΔE = g.tryRotation(row_storage);

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

    if (accepted) {
      E += ΔE;
      g.acceptRotation(row_storage);
    }

    if (!dry) {
      A.after_step(accepted, g, θ0, E, ΔE, M);
    }
    return accepted;
  }

  double sweep(bool dry = false) {
    unsigned total = 0;
    unsigned accepted = 0;

    for (unsigned i = 0; i < M.size() - 1; i++) {
      for (unsigned j = i + 1; j < M.size(); j++) {
        Givens g1(M, false, i, j, θ0, rng);
        Givens g2(M, true, i, j, θ0, rng);

        if (this->step(g1, dry))
          accepted++;
        if (this->step(g2, dry))
          accepted++;

        total += 2;
      }
    }

    return (double)accepted / (double)total;
  }

  void tune(unsigned N, double ε) {
    for (unsigned i = 0; i < N; i++) {
      double ratio_accepted = this->sweep(true);
      if (ratio_accepted > 0.5) {
        θ0 *= 1 + ε;
      } else {
        θ0 /= 1 + ε;
      }
    }
  }

  double run(unsigned N, bool dry = false, unsigned nOrth = 1000) {
    double totalRatio = 0;

    for (unsigned i = 0; i < N; i++) {
      totalRatio += this->sweep(dry);
      if (!dry) {
        A.after_sweep(θ0, E, M);
      }
      if (i % nOrth == 0) {
        M.orthogonalize();
      }
    }

    return totalRatio / N;
  }
};