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

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

  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 double
  integrate(const Polynomial<N>& P, const double& xinf, const double& xsup)
  {
    Polynomial<N + 1> Q = primitive(P);
    return (Q(xsup) - Q(xinf));
  }

  PUGS_INLINE
  constexpr friend 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);
    }
  }

  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
  lagrangePolynomial(const TinyVector<N + 1> zeros, const size_t k, Polynomial<N>& lk)
  {
    for (size_t i = 0; i < N; ++i) {
      Assert(zeros[i] < zeros[i + 1], "Interpolation values must be strictly increasing in Lagrange polynomials");
    }
    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;
    }
  }

  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;
};
#endif   // POLYNOMIAL_HPP