#ifndef POLYNOMIAL_HPP
#define POLYNOMIAL_HPP

#include <algebra/TinyVector.hpp>

template <size_t N>
class Polynomial
{
 private:
  using Coefficients = TinyVector<N + 1, double>;
  Coefficients m_coefficients;
  static_assert((N >= 0), "Polynomial degree must be non-negative");

 public:
  PUGS_INLINE
  constexpr size_t
  degree() const
  {
    return N;
  }
  PUGS_INLINE
  constexpr size_t
  realDegree() const
  {
    for (size_t j = N; j > 0; j--) {
      if (std::abs(coefficients()[j]) > 1.e-14) {
        return j;
      }
    }
    return 0;
  }
  PUGS_INLINE
  constexpr const TinyVector<N + 1>&
  coefficients() const
  {
    return m_coefficients;
  }

  PUGS_INLINE
  constexpr TinyVector<N + 1>&
  coefficients()
  {
    return m_coefficients;
  }

  template <size_t M>
  PUGS_INLINE constexpr bool
  operator==(const Polynomial<M>& q) const
  {
    if (M != N)
      return false;
    if (m_coefficients != q.coefficients()) {
      return false;
    }
    return true;
  }

  PUGS_INLINE
  constexpr bool
  operator!=(const Polynomial<N>& q) const
  {
    return not this->operator==(q);
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<std::max(M, N)>
  operator+(const Polynomial<M>& Q) const
  {
    Polynomial<std::max(M, N)> P;
    if constexpr (M > N) {
      P.coefficients() = Q.coefficients();
      for (size_t i = 0; i <= N; ++i) {
        P.coefficients()[i] += coefficients()[i];
      }
    } else {
      P.coefficients() = coefficients();
      for (size_t i = 0; i <= M; ++i) {
        P.coefficients()[i] += Q.coefficients()[i];
      }
    }
    return P;
  }

  PUGS_INLINE constexpr Polynomial<N>
  operator-() const
  {
    Polynomial<N> P;
    P.coefficients() = -coefficients();
    return P;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<std::max(M, N)>
  operator-(const Polynomial<M>& Q) const
  {
    Polynomial<std::max(M, N)> P;
    if constexpr (M > N) {
      P.coefficients() = -Q.coefficients();
      for (size_t i = 0; i <= N; ++i) {
        P.coefficients()[i] += coefficients()[i];
      }
    } else {
      P.coefficients() = coefficients();
      for (size_t i = 0; i <= M; ++i) {
        P.coefficients()[i] -= Q.coefficients()[i];
      }
    }
    return P;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<N>&
  operator=(const Polynomial<M>& Q)
  {
    coefficients() = zero;
    for (size_t i = N + 1; i <= M; ++i) {
      Assert(Q.coefficients()[i] == 0, "degree of polynomial to small in assignation");
    }
    //    static_assert(N >= M, "degree of polynomial to small in assignation");
    for (size_t i = 0; i <= std::min(M, N); ++i) {
      coefficients()[i] = Q.coefficients()[i];
    }
    return *this;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<N>&
  operator+=(const Polynomial<M>& Q)
  {
    static_assert(N >= M, "Polynomial degree to small in affectation addition");
    for (size_t i = 0; i <= M; ++i) {
      coefficients()[i] += Q.coefficients()[i];
    }
    return *this;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<N>&
  operator-=(const Polynomial<M>& Q)
  {
    static_assert(N >= M, "Polynomial degree to small in affectation addition");
    for (size_t i = 0; i <= M; ++i) {
      coefficients()[i] -= Q.coefficients()[i];
    }
    return *this;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<M + N> operator*(const Polynomial<M>& Q) const
  {
    Polynomial<M + N> P;
    P.coefficients() = zero;
    for (size_t i = 0; i <= N; ++i) {
      for (size_t j = 0; j <= M; ++j) {
        P.coefficients()[i + j] += coefficients()[i] * Q.coefficients()[j];
      }
    }
    return P;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<N>&
  operator*=(const Polynomial<M>& Q)
  {
    static_assert(N >= M, "Degree to small in affectation product");
    for (size_t i = N - M + 1; i <= N; ++i) {
      Assert(coefficients()[i] == 0, "Degree of affectation product greater than the degree of input polynomial");
    }
    Polynomial<N> P(zero);
    for (size_t i = 0; i <= N - M; ++i) {
      for (size_t j = 0; j <= M; ++j) {
        P.coefficients()[i + j] += coefficients()[i] * Q.coefficients()[j];
      }
    }
    coefficients() = P.coefficients();
    return *this;
  }

  PUGS_INLINE
  constexpr Polynomial<N> operator*(const double& lambda) const
  {
    TinyVector<N + 1> mult_coefs = lambda * m_coefficients;
    Polynomial<N> M(mult_coefs);
    return M;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<M * N>
  compose(const Polynomial<M>& Q) const
  {
    Polynomial<M * N> P;
    P.coefficients() = zero;
    Polynomial<M * N> R;
    R.coefficients() = zero;

    for (size_t i = 0; i <= N; ++i) {
      R = Q.template pow<N>(i) * coefficients()[i];
      P += R;   // R;
    }
    return P;
  }

  template <size_t M, size_t I>
  PUGS_INLINE constexpr Polynomial<M * N>
  power(const Polynomial<M>& Q) const
  {
    return coefficients()[I] * Q.template pow<N>(I);
  }

  template <size_t M, size_t... I>
  PUGS_INLINE constexpr Polynomial<M * N>
  compose2(const Polynomial<M>& Q, std::index_sequence<I...>) const
  {
    Polynomial<M * N> P;
    P.coefficients() = zero;
    P                = (power<M, I>(Q) + ...);
    return P;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<M * N>
  compose2(const Polynomial<M>& Q) const
  {
    Polynomial<M * N> P;
    P.coefficients()    = zero;
    using IndexSequence = std::make_index_sequence<N + 1>;
    return compose2<M>(Q, IndexSequence{});
    //    for (size_t i = 0; i <= N; ++i) {
    //      P += Q.template pow<N>(i) * coefficients()[i];
    //    }
    return P;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<M * N>
  operator()(const Polynomial<M>& Q) const
  {
    Polynomial<M * N> P;
    P.coefficients() = zero;
    Polynomial<M * N> R;
    R.coefficients() = zero;

    for (size_t i = 0; i <= N; ++i) {
      R = Q.template pow<N>(i) * coefficients()[i];
      P += R;   // R;
    }
    return P;
  }

  template <size_t M>
  PUGS_INLINE constexpr Polynomial<M * N>
  pow(size_t power) const
  {
    Assert(power <= M, "You declared a polynomial of degree too small for return of the pow function");
    Polynomial<M * N> R;
    R.coefficients() = zero;
    if (power == 0) {
      R.coefficients()[0] = 1;
    } else {
      R = *this;
      for (size_t i = 1; i < power; ++i) {
        R = R * *this;
      }
    }
    return R;
  }

  PUGS_INLINE
  constexpr friend Polynomial<N> operator*(const double& lambda, const Polynomial<N> P)
  {
    return P * lambda;
  }
  // evaluation using Horner's method https://en.wikipedia.org/wiki/Horner's_method
  PUGS_INLINE
  constexpr double
  evaluate(const double& x) const
  {
    TinyVector<N + 1> coefs = this->coefficients();
    double bcoef            = coefs[N];
    for (size_t i = N; i > 0; --i) {
      bcoef *= x;
      bcoef += coefs[i - 1];
    }
    return bcoef;
  }

  PUGS_INLINE
  constexpr double
  operator()(const double x) const
  {
    TinyVector<N + 1> coefs = this->coefficients();
    double bcoef            = coefs[N];
    for (size_t i = N; i > 0; --i) {
      bcoef *= x;
      bcoef += coefs[i - 1];
    }
    return bcoef;
  }

  template <size_t M>
  PUGS_INLINE constexpr friend void
  divide(const Polynomial<N>& P1, const Polynomial<M>& P2, Polynomial<N>& Q, Polynomial<N>& R)
  {
    const size_t Nr  = P1.realDegree();
    const size_t Mr  = P2.realDegree();
    R.coefficients() = P1.coefficients();
    Q.coefficients() = zero;
    for (size_t k = Nr - Mr + 1; k > 0; --k) {
      Q.coefficients()[k - 1] = R.coefficients()[Mr + k - 1] / P2.coefficients()[Mr];
      for (size_t j = Mr + k - 1; j > (k - 1); --j) {
        R.coefficients()[j] -= Q.coefficients()[k - 1] * P2.coefficients()[j - k];
      }
    }
    for (size_t j = Mr; j < Nr + 1; ++j) {
      R.coefficients()[j] = 0;
    }
  }

  PUGS_INLINE
  constexpr friend Polynomial<N + 1>
  primitive(const Polynomial<N>& P)
  {
    TinyVector<N + 2> coefs;
    for (size_t i = 0; i < N + 1; ++i) {
      coefs[i + 1] = P.coefficients()[i] / double(i + 1);
    }
    coefs[0] = 0;
    return Polynomial<N + 1>{coefs};
  }

  PUGS_INLINE
  constexpr friend std::ostream&
  operator<<(std::ostream& os, const Polynomial<N>& P)
  {
    //    os << "P(x) = ";
    bool all_coef_zero = true;
    if (N == 0) {
      os << P.coefficients()[0];
      return os;
    }
    if (N != 1) {
      if (P.coefficients()[N] != 0.) {
        if (P.coefficients()[N] < 0.) {
          os << "- ";
        }
        if (P.coefficients()[N] != 1 && P.coefficients()[N] != -1) {
          os << std::abs(P.coefficients()[N]);
        }
        os << "x^" << N;
        all_coef_zero = false;
      }
    }
    for (size_t i = N - 1; i > 1; --i) {
      if (P.coefficients()[i] != 0.) {
        if (P.coefficients()[i] > 0.) {
          os << " + ";
        } else if (P.coefficients()[i] < 0.) {
          os << " - ";
        }
        if (P.coefficients()[i] != 1 && P.coefficients()[i] != -1) {
          os << std::abs(P.coefficients()[i]);
        }
        os << "x^" << i;
        all_coef_zero = false;
      }
    }
    if (P.coefficients()[1] != 0.) {
      if (P.coefficients()[1] > 0. && N != 1) {
        os << " + ";
      } else if (P.coefficients()[1] < 0.) {
        os << " - ";
      }
      if (P.coefficients()[1] != 1 && P.coefficients()[1] != -1) {
        os << std::abs(P.coefficients()[1]);
      }
      os << "x";
      all_coef_zero = false;
    }
    if (P.coefficients()[0] != 0. || all_coef_zero) {
      if (P.coefficients()[0] > 0.) {
        os << " + ";
      } else if (P.coefficients()[0] < 0.) {
        os << " - ";
      }
      os << std::abs(P.coefficients()[0]);
    }
    return os;
  }

  PUGS_INLINE
  constexpr friend void
  lagrangeBasis(const TinyVector<N + 1> zeros, TinyVector<N + 1, Polynomial<N>>& basis)
  {
    Polynomial<N> lj;
    for (size_t j = 0; j < N + 1; ++j) {
      basis[j] = lagrangePolynomial(zeros, j);
    }
  }

  PUGS_INLINE
  constexpr friend Polynomial<N>
  lagrangeToCanonical(const TinyVector<N + 1> lagrange_coefs, const TinyVector<N + 1, Polynomial<N>>& basis)
  {
    Polynomial<N> P(zero);
    //    lagrangeBasis({0, 0, 0}, basis);
    for (size_t j = 0; j < N + 1; ++j) {
      P += basis[j] * lagrange_coefs[j];
    }
    return P;
  }

  PUGS_INLINE constexpr Polynomial(const TinyVector<N + 1>& coefficients) noexcept : m_coefficients{coefficients} {}

  PUGS_INLINE
  constexpr Polynomial(TinyVector<N + 1>&& coefficients) noexcept : m_coefficients{coefficients} {}

  PUGS_INLINE
  constexpr Polynomial() noexcept = default;
  ~Polynomial()                   = default;
};

template <size_t N>
PUGS_INLINE constexpr Polynomial<N> lagrangePolynomial(const TinyVector<N + 1> zeros, const size_t k);

template <size_t N>
PUGS_INLINE constexpr TinyVector<N, Polynomial<N - 1>>
lagrangeBasis(const TinyVector<N>& zeros)
{
  static_assert(N >= 1, "invalid degree");
  TinyVector<N, Polynomial<N - 1>> basis;
  Polynomial<N - 1> lj;
  for (size_t j = 0; j < N; ++j) {
    basis[j] = lagrangePolynomial<N - 1>(zeros, j);
  }
  return basis;
}

template <size_t N>
PUGS_INLINE constexpr double
integrate(const Polynomial<N>& P, const double& xinf, const double& xsup)
{
  Polynomial<N + 1> Q = primitive(P);
  return (Q(xsup) - Q(xinf));
}

template <size_t N>
PUGS_INLINE constexpr double
symmetricIntegrate(const Polynomial<N>& P, const double& delta)
{
  Assert(delta > 0, "A positive delta is needed for symmetricIntegrate");
  double integral = 0.;
  for (size_t i = 0; i <= N; ++i) {
    if (i % 2 == 0)
      integral += 2. * P.coefficients()[i] * std::pow(delta, i + 1) / (i + 1);
  }
  return integral;
}

template <size_t N>
PUGS_INLINE constexpr auto
derivative(const Polynomial<N>& P)
{
  if constexpr (N == 0) {
    return Polynomial<0>(0);
  } else {
    TinyVector<N> coefs;
    for (size_t i = 0; i < N; ++i) {
      coefs[i] = double(i + 1) * P.coefficients()[i + 1];
    }
    return Polynomial<N - 1>(coefs);
  }
}

template <size_t N>
PUGS_INLINE constexpr Polynomial<N>
lagrangePolynomial(const TinyVector<N + 1> zeros, const size_t k)
{
  for (size_t i = 0; i < N; ++i) {
    Assert(zeros[i] < zeros[i + 1], "Interpolation values must be strictly increasing in Lagrange polynomials");
  }
  Polynomial<N> lk;
  lk.coefficients()    = zero;
  lk.coefficients()[0] = 1;
  for (size_t i = 0; i < N + 1; ++i) {
    if (k == i)
      continue;
    double factor = 1. / (zeros[k] - zeros[i]);
    Polynomial<1> P({-zeros[i] * factor, factor});
    lk *= P;
  }
  return lk;
}

#endif   // POLYNOMIAL_HPP