#include "space_wolff.hpp"

template <int D> using Sphere = Spin<double, D, Radius>;

template <int D> double softPotential(double a, double k, const Vector<double, D>& diff, double σ) {
  double δ = σ - sqrt(diff.transpose() * diff);

  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> double hardPotential(const Vector<double, D>& diff, double σ) {
  if (pow(σ, 2) < diff.transpose() * diff) {
    return 0;
  } else {
    return -std::numeric_limits<double>::infinity();
  }
}

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;
    return softPotential(a, k, d, s1.s + s2.s);
  };
}

template <int D>
std::function<double(const Spin<double, D, Radius>&, const Spin<double, D, Radius>&)>
zSpheresTorus(double L, double a, double k) {
  return [L, a, k](const Spin<double, D, double>& s1, const Spin<double, D, double>& s2) -> double {
    return softPotential(a, k, diff(L, s1.x, s2.x), s1.s + s2.s);
  };
}

template <int D>
std::function<double(const Spin<double, D, Dimer<double, D>>&,
                     const Spin<double, D, Dimer<double, D>>&)>
zDimers(std::function<double(const Sphere<D>&, const Sphere<D>&)> zSingle) {
  return [zSingle](const Spin<double, D, Dimer<double, D>>& s1,
            const Spin<double, D, Dimer<double, D>>& s2) -> double {
    Spin<double, D, Radius> s11 = {.x = s1.x + s1.s.relativePosition, .s = s1.s.radius};
    Spin<double, D, Radius> s12 = {.x = s1.x - s1.s.relativePosition, .s = s1.s.radius};
    Spin<double, D, Radius> s21 = {.x = s2.x + s2.s.relativePosition, .s = s2.s.radius};
    Spin<double, D, Radius> s22 = {.x = s2.x - s2.s.relativePosition, .s = s2.s.radius};

    return zSingle(s11, s21) + zSingle(s12, s21) + zSingle(s11, s22) + zSingle(s12, s22);
  };
}

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);
  };
}