#include "space_wolff.hpp"

template <int D>
std::function<double(const Spin<double, D, Radius>&, const Spin<double, D, Radius>&)>
zSpheres(double a, double k) {
  return [a, k](const Spin<double, D, Radius>& s1, const Spin<double, D, Radius>& s2) -> double {
    Vector<double, D> d = s1.x - s2.x;

    double σ = s1.s + s2.s;
    double δ = σ - sqrt(d.transpose() * d);

    if (δ > -a * σ) {
      return 0.5 * k * (2 * pow(a * σ, 2) - pow(δ, 2));
    } else if (δ > -2 * a * σ) {
      return 0.5 * k * pow(δ + 2 * a * σ, 2);
    } else {
      return 0;
    }
  };
}

template <int D, class S> std::function<double(Spin<double, D, S>)> bCenter(double H) {
  return [H](Spin<double, D, S> s) -> double { return H * s.x.norm(); };
}

template <int D> Euclidean<double, D> flipAround(double θ, Vector<double, D>& t₀) {
  Matrix<double, D> m;

  m(0, 0) = -cos(2 * θ);
  m(1, 1) = cos(2 * θ);
  m(0, 1) = -2 * cos(θ) * sin(θ);
  m(1, 0) = -2 * cos(θ) * sin(θ);

  return Euclidean<double, D>(t₀ - m * t₀, m);
}

template <int D, class S> Gen<double, D, Euclidean<double, D>, S> nudgeGen(double ε) {
  return [ε](Model<double, D, Euclidean<double, D>, S>& M,
             Rng& rng) -> Transformation<double, D, Euclidean<double, D>, S>* {
    double θ = rng.uniform((double)0.0, 2 * M_PI);

    Spin<double, D, S>* s = rng.pick(M.s);
    Vector<double, D> t = s->x;
    for (unsigned j = 0; j < D; j++) {
      t(j) += rng.variate<double, std::normal_distribution>(0.0, ε);
    }

    return new SpinFlip<double, D, Euclidean<double, D>, S>(M, flipAround(θ, t), s);
  };
}

template <int D, class S> Gen<double, D, Euclidean<double, D>, S> swapGen(double ε) {
  return [ε](Model<double, D, Euclidean<double, D>, S>& M,
             Rng& rng) -> Transformation<double, D, Euclidean<double, D>, S>* {
    Spin<double, D, S>* s1 = rng.pick(M.s);
    Spin<double, D, S>* s2 = rng.pick(M.s);

    while (s1 == s2) {
      s2 = rng.pick(M.s);
    }

    Vector<double, D> t1 = s1->x;
    Vector<double, D> t2 = s2->x;
    Vector<double, D> t12 = t1 - t2;
    Vector<double, D> t = (t1 + t2) / 2;

    double θ =
        atan2(t12(1), t12(0)) + rng.variate<double, std::normal_distribution>(0.0, ε) / t12.norm();

    return new PairFlip<double, D, Euclidean<double, D>, S>(M, flipAround(θ, t), s1, s2);
  };
}

template <int D, class S> Gen<double, D, Euclidean<double, D>, S> accrossGen(double ε) {
  return [ε](Model<double, D, Euclidean<double, D>, S>& M,
             Rng& rng) -> Transformation<double, D, Euclidean<double, D>, S>* {
    Spin<double, D, S>* s1 = rng.pick(M.s);
    Spin<double, D, S>* s2 = rng.pick(M.s);

    while (s1 == s2) {
      s2 = rng.pick(M.s);
    }

    Vector<double, D> t1 = s1->x;
    Vector<double, D> t2 = s2->x;
    Vector<double, D> t12 = t1 - t2;
    Vector<double, D> t = t2;

    double θ =
        atan2(t12(1), t12(0)) + rng.variate<double, std::normal_distribution>(0.0, ε) / t12.norm();

    return new SpinFlip<double, D, Euclidean<double, D>, S>(M, flipAround(θ, t), s1);
  };
}

template <int D, class S> Gen<double, D, Euclidean<double, D>, S> centerGen(double ε) {
  return [ε](Model<double, D, Euclidean<double, D>, S>& M,
             Rng& rng) -> Transformation<double, D, Euclidean<double, D>, S>* {
    double θ = rng.uniform((double)0.0, 2 * M_PI);

    Vector<double, D> t = M.s0.t;
    for (unsigned j = 0; j < D; j++) {
      t(j) += rng.variate<double, std::normal_distribution>(0.0, ε);
    }

    Spin<double, D, S>* s = rng.pick(M.s);

    return new SpinFlip<double, D, Euclidean<double, D>, S>(M, flipAround(θ, t), s);
  };
}