#pragma once
#include "config.hpp" // FUNC_USE_BOOST
#include <cmath> // sqrt
#include <limits> // std::numeric_limits<T>::digits
#include <array> // std::array
#include <utility> // std::pair
#include <type_traits>

#include "FunctionContainer.hpp"

#ifdef FUNC_USE_BOOST
#define BOOST_MATH_GAUSS_NO_COMPUTE_ON_DEMAND
#include <boost/math/quadrature/gauss_kronrod.hpp> // gauss_kronrod::integrate
#include <boost/math/tools/roots.hpp> // newton_raphson_iterate
#endif

namespace func {


/**
  \brief Transforms a uniformly spaced partition of \f$[a,b]\f$ into a nonuniform partition of \f$[a,b]\f$.

  For efficiency, we require a Transfer function is simply an increasing cubic
  polynomial such that \f$p(a)=0, p(b)=b/\mathrm{stepSize}\f$.  To help compete
  against uniform lookup tables, part of the `operator()` must be baked into
  those coefficients (explaining why stepSize appears in the formula above).

  When given a FunctionContainer with a defined first derivative then this will construct a valid TransferFunction.
  Let \f$f\f$ be a function defined over \f$[a,b]\f$. Then, construct \f$S: [a,b]\rightarrow[a,b]\f$ as
    \f[S(x) = a + \frac{b-a}{c}\int_a^x\frac{1}{\sqrt{1+f'(t)^2}} dt,\f]
  where \f$c = \int_a^b\frac{1}{\sqrt{1+f'(t)^2}} dt\f$. This way, \f$S(a)=a, S(b)=b\f$.
  Computing \f$S^{-1}\f$ has to be fast so we approximate is as a monotone Hermite
  cubic polynomial and then rebuild \f$S\f$ as \f$(S^{-1})^{-1}\f$ using
  Boost's newton_raphson_iterate.

  \note Can be constructed with a std::array<TIN,4> of polynomial coefficients.
  \note An example: the identity transfer function which is the linear polynomial with coefs {-m_minArg/m_stepSize,1/m_stepSize,0,0}.

  \note We need Boost version 1.71.0 or newer to generate a candidate transfer function.
    Boost is not required if TransferFunction is given coefficients.
  \note Resulting nonuniform LUT performs better if f' is largest near the endpoints [a,b].
    If f' is largest near the middle of an interval (eg \f$e^{-x^2}\f$ when \f$a<-3\f$ and \f$b>3\f$)
    then the grid remains uniform. This is an issue with the way we approximate S!

  \todo Currently, transfer functions are basically useless if \f$f'\f$ is not
  extreme near the endpoints of its domain. (b/c we use cubic hermite
  interpolation at the endpoints).

  We want to accurately approximate \f$S\f$ anywhere \f$f'(x)\f$ is extreme!
  The best possible polynomial approximating \f$S\f$ might be promising...
  
  \note A cubic polynomial
    \f$p(x) = a_0 + a_1x + a_2x^2 + a_3x^3\f$
  is monotone over \f$\mathbb{R}\f$ if and only if \f$a_2^2 < 3a_1a_3\f$. Maybe
  we can compute a polynomial minimizing
    \f$\max_t |f(t)-p(t)| \f$
  such that
    \f$p(a) = a, p(b) = b\f$, and \f$a_2^2 < 3a_1a_3\f$.
  That would be a much better general purpose solution. The search space is convex!
 */
template <typename TIN>
class TransferFunction
{
  TIN m_minArg, m_tableMaxArg, m_stepSize;

  /** This class is a polynomial approximating g inverse.
   * The identity transfer function is {-m_minArg/m_stepSize,1/m_stepSize,0,0} */
  __attribute__((aligned)) std::array<TIN,4> m_inverse_coefs = {{0,0,0,0}};
public:
  /* Set m_inverse_coefs equal to a vector that came from a json file (or TODO from a LookupTableParameters object) */
  TransferFunction(const std::array<TIN,4>& inv_coefs) { m_inverse_coefs = inv_coefs; }
  //TransferFunction(const TransferFunction<TIN>& transferFunction) :
  //  m_inverse_coefs(transferFunction.m_inverse_coefs), m_minArg(transferFunction.m_minArg),
  //  m_tableMaxArg(transferFunction.m_tableMaxArg), m_stepSize(transferFunction.m_stepSize) {}
  //
  //TransferFunction<TIN>& operator=(TransferFunction<TIN> transferFunction){
  //  m_inverse_coefs = transferFunction.m_inverse_coefs;
  //  m_minArg = transferFunction.m_minArg;
  //  m_tableMaxArg = transferFunction.m_tableMaxArg;
  //  m_stepSize = transferFunction.m_stepSize;
  //  return *this;
  //}

  TransferFunction() = default;

  /** Build the coefficients in g_inv */
  template<typename TOUT>
  TransferFunction(const FunctionContainer<TIN,TOUT>& fc, TIN minArg, TIN tableMaxArg, TIN stepSize) : 
    m_minArg(minArg), m_tableMaxArg(tableMaxArg), m_stepSize(stepSize) {
#ifndef FUNC_USE_BOOST
    throw std::invalid_argument("func::TransferFunction cannot generate a transfer function without boost installed.");
#else
    using boost::math::quadrature::gauss_kronrod;
    using boost::math::differentiation::make_fvar;

    /* TODO can we use SFINAE to check if sqrt(TOUT) -> TIN is defined? */

    auto fun = fc.autodiff1_fun;
    if(fun == nullptr)
      throw std::invalid_argument("Error in func::TransferFunction. 1st derivative of function is needed to generate nonuniform grids but is null");

    /* std::function to return the first derivative of f */
    std::function<TOUT(TIN)> f_prime = [&fun](TIN x) -> TOUT {
      return fun(make_fvar<TIN,1>(x)).derivative(1);
    };

    /* build the integrand. Notice that it is strictly positive */
    std::function<TIN(TIN)> integrand = [&f_prime](TIN x) -> TIN {
      return static_cast<TIN>(1.0)/static_cast<TIN>(sqrt(static_cast<TOUT>(1.0) + f_prime(x)*f_prime(x)));
    };

    TIN a = m_minArg;
    TIN b = m_tableMaxArg;

    /* integrate over [a,b] using adaptive quadrature & default tol of sqrt(epsilon_machine). */
    TIN c = gauss_kronrod<TIN, 15>::integrate(integrand, a, b);

    /* rescale the integrand to map [a,b] -> [a,b] */
    std::function<TIN(TIN)> g_prime = [&integrand,a,b,c](TIN x) -> TIN {
      return (b-a) * integrand(x)/c;
    };

    /* g(a)=a and g(b)=b so
     * g_inv_prime(a) = 1/g_prime(a) and g_inv_prime(b) = 1/g_prime(b) */
    TIN m0 = 1/g_prime(a);
    TIN m1 = 1/g_prime(b);

    /* Ensure monotonicity of the Hermite interpolating polynomial. Note that m0,m1 >= 0.
     * TODO this is sufficient to ensure p is monotone, but it's not necessary. Figure out how to allow m0>3 or m1>3. */
    m0 = (m0 > 3) ? 3 : m0;
    m1 = (m1 > 3) ? 3 : m1;

    /* Compute the Hermite interpolating polynomial p for g_inv satisfying
     * p(a)=a, p(b)=b, p'(a)=m0, p'(b)=m1
     * (this symbolic expression was computed by Matlab) */
    m_inverse_coefs[0] = (a*b*(a + b - a*m1 - b*m0))/(a - b)/(a - b);
    m_inverse_coefs[1] = (a*a*m1 - 6*a*b + b*b*m0 + 2*a*b*m0 + 2*a*b*m1)/(a - b)/(a - b);
    m_inverse_coefs[2] = -(a*m0 - 3*b - 3*a + 2*a*m1 + 2*b*m0 + b*m1)/(a - b)/(a - b);
    m_inverse_coefs[3] = (m0 + m1 - 2)/(a - b)/(a - b);

    /* build the inverse function that we'll actually use by encoding the underlying table's hash into
     * the transfer function eval. This way, the only cost of using a transfer function is 4 stack IOPs and 6 FLOPs */
    m_inverse_coefs[0] = m_inverse_coefs[0] - m_minArg;
    for(unsigned int i=0; i<4; i++)
      m_inverse_coefs[i] = m_inverse_coefs[i] / m_stepSize;
#endif
  }

  /** Evaluate the inverse using Horner's method */
  TIN inverse(TIN x) const {
    TIN sum = static_cast<TIN>(0);
    for (int k=3; k>0; k--)
      sum = x*(m_inverse_coefs[k] + sum);
    return sum + m_inverse_coefs[0];
  }

  /** Evaluate the derivative of the inverse using Horner's method */
  TIN inverse_diff(TIN x) const {
    TIN sum = static_cast<TIN>(0);
    for (int k=3; k>1; k--)
      sum = x*(k*m_inverse_coefs[k] + sum);
    return sum + m_inverse_coefs[1];
  }

  /** Use newton-raphson_iteration on p. Recall that each coef in
   * g was divided by h and we subtracted by m_minArg */
  TIN operator()(TIN x) const {
    // This will have at least 0.9*std::numeric_limits<TIN>::digits digits of accuracy
    boost::uintmax_t maxit = 55;
    return boost::math::tools::newton_raphson_iterate(
        [this,x](const TIN z){ return std::make_tuple(inverse(z) - (x-m_minArg)/m_stepSize, inverse_diff(z));},
        x,m_minArg, m_tableMaxArg, 0.9*std::numeric_limits<TIN>::digits, maxit);
  }

  std::array<TIN,4> get_coefs() const { return m_inverse_coefs; }
  TIN min_arg() const { return m_minArg; }
  TIN max_arg() const { return m_tableMaxArg; }
};

/** print basic info about a TransferFunction */
template <typename TIN>
std::ostream& operator<<(std::ostream& out, const TransferFunction<TIN>& F){
  auto coefs = F.get_coefs();
  out << "degree 3 monotone Hermite interpolation with polynomial: \n";
  out << std::to_string(coefs[3]) << "x^3 + " << std::to_string(coefs[2]) << "x^2 + " <<
    std::to_string(coefs[1]) << "x + " << std::to_string(coefs[0]) << ". Defined over [" << F.min_arg() << "," << F.max_arg() << "].\n";
  return out;
}

/** wraps operator<< */
template <typename TIN>
inline std::string to_string(const TransferFunction<TIN>& F) {
  std::ostringstream ss;
  ss << F;
  return ss.str();
}

}
