#pragma once

#include "mod.hpp"
#include "matrix.hpp"
#include "spin.hpp"
#include "vector.hpp"

typedef double Radius;
typedef signed IsingSpin;

template <class T, int D> class Dimer {
public:
  double radius;
  Vector<T, D> relativePosition;
};

template <class T, int D> class Affine {
protected:
  virtual Vector<T, D> actV(const Vector<T, D>& x) const { return t + r * x; }
  Affine actA(const Affine& x) const { return Affine(act(x.t), r * x.r); }

public:
  Vector<T, D> t;
  Matrix<T, D> r;

  Affine() {
    for (unsigned i = 0; i < D; i++) {
      t(i) = 0;
      r(i, i) = 1;
      for (unsigned j = 1; j < D; j++) {
        r(i, (i + j) % D) = 0;
      }
    }
  }

  Affine(const Vector<T, D>& t0, const Matrix<T, D>& r0) : t(t0), r(r0) {}

  Affine inverse() const { return Affine(-r.transpose() * t, r.transpose()); }

  Vector<T, D> act(const Vector<T, D>& x) const { return actV(x); }
  Radius act(Radius r) const { return r; }
  IsingSpin act(IsingSpin s) const { return s; }
  Dimer<T, D> act(const Dimer<T, D>& d) const {
    return {.radius = d.radius, .relativePosition = r * d.relativePosition};
  }

  template <class S> Spin<T, D, S> act(const Spin<T, D, S>& s) const {
    return {.x = act(s.x), .s = act(s.s)};
  }
};

template <class T, int D> class Euclidean : public Affine<T, D> {
public:
  Euclidean(T L = 0) : Affine<T, D>() {}
  Euclidean(const Affine<T, D>& a) : Affine<T, D>(a) {}
  Euclidean(Vector<T, D> t0, Matrix<T, D> r0) : Affine<T, D>(t0, r0) {}
  
  Euclidean act(const Euclidean& t) const { return Euclidean(Affine<T, D>::actA(t)); };
  Euclidean inverse() const { return Euclidean(Affine<T, D>::inverse()); }

  using Affine<T, D>::act;
};

template <class T, int D> class TorusGroup : public Affine<T, D>{
private:
  T L;

protected:
  Vector<T, D> actV(const Vector<T, D>& s) const override {
    Vector<T, D> s_new = Affine<T, D>::actV(s);

    for (unsigned i = 0; i < D; i++) {
      s_new(i) = mod(s_new(i), L);
    }

    return s_new;
  }

public:
  TorusGroup(T L) : Affine<T, D>(), L(L) {}
  TorusGroup(T L, const Affine<T, D>& t) : Affine<T, D>(t), L(L) {}
  TorusGroup(T L, Vector<T, D> t0, Matrix<T, D> r0) : Affine<T, D>(t0, r0), L(L) {}

  TorusGroup act(const TorusGroup& t) const { return TorusGroup(L, Affine<T, D>::actA(t)); }
  TorusGroup inverse() const { return TorusGroup(L, Affine<T, D>::inverse()); }

  using Affine<T, D>::act;
};