#ifndef _utl_RK5ODEIntegrator_h_
#define _utl_RK5ODEIntegrator_h_

#include <cmath>
#include <utl/SafeBoolCast.h>


namespace utl {

  /**
    \class RK5Iterator RK5ODEIntegrator.h "utl/RK5ODEIntegrator.h"

    This iterator class will advance the solution for the current point.
    See class RK5ODEIntegrator for details.

    \author Darko Veberic
    \date 20 Jun 2008
    \version $Id$

    \ingroup math
  */
  template<class DerivativeFunctor, class Vector>
  class RK5Iterator {
  public:
    RK5Iterator(const double x, const Vector& y,
                DerivativeFunctor& d)
      : fX(x), fDerivativeFunctor(d)
    {
      Copy(d, y, fY);
      Clear(d, fYError);
    }

    double GetX() const { return fX; }

    Vector& GetY() { return fY; }

    const Vector& GetY() const { return fY; }

    const Vector& GetYError() const { return fYError; }

    void
    Advance(const double dx, const Vector& dYdX)
    {
      Step(dx, fY, dYdX);
      fX += dx;
    }

    void
    Advance(const double dx)
    {
      Vector dYdX;
      fDerivativeFunctor(fX, fY, dYdX);
      Advance(dx, dYdX);
    }

    RK5Iterator&
    operator+=(const double dx)
    { Advance(dx); return *this; }

    RK5Iterator&
    operator-=(const double dx)
    { Advance(-dx); return *this; }

    DerivativeFunctor& GetDerivativeFunctor()
    { return fDerivativeFunctor; }

  protected:
    void
    Step(const double dx, Vector& yNew, const Vector& dYdX)
    {
      const unsigned int n = fDerivativeFunctor;
      Vector y1;
      Vector dYdXt[5];
      for (unsigned int k = 0; k < 5; ++k) {
        for (unsigned int i = 0; i < n; ++i) {
          double dy = fB[k][0] * dYdX[i];
          for (unsigned int j = 1; j <= k; ++j)
            dy += fB[k][j] * dYdXt[j-1][i];
          y1[i] = fY[i] + dx * dy;
        }
        fDerivativeFunctor(fX + fA[k] * dx, y1, dYdXt[k]);
      }
      for (unsigned int i = 0; i < n; ++i) {
        double dy = fC[0] * dYdX[i];
        double de = fD[0] * dYdX[i];
        for (unsigned int j = 1; j < 5; ++j) {
          dy += fC[j] * dYdXt[j][i];
          de += fD[j] * dYdXt[j][i];
        }
        yNew[i] = fY[i] + dx * dy;
        fYError[i] = dx * de;
      }
    }

    static
    void
    Copy(const unsigned int n, const Vector& a, Vector& b)
    {
      for (unsigned int i = 0; i < n; ++i)
        b[i] = a[i];
    }

    static
    void
    Clear(const unsigned int n, Vector& b)
    {
      for (unsigned int i = 0; i < n; ++i)
        b[i] = 0;
    }

    double fX;
    Vector fY;
    Vector fYError;
    DerivativeFunctor& fDerivativeFunctor;

    // constants
    static const double fA[6];
    static const double fB[5][5];
    static const double fC[5];
    static const double fD[5];
  };

  template<class DerivativeFunctor, class Vector>
  const double RK5Iterator<DerivativeFunctor, Vector>::fA[] = {
    0.2, 0.3, 0.6, 1, 0.875
  };

  template<class DerivativeFunctor, class Vector>
  const double RK5Iterator<DerivativeFunctor, Vector>::fB[][5] = {
    { 0.2 },
    { 0.075, 0.225 },
    { 0.3, -0.9, 1.2 },
    { -11./54, 2.5, -70./27, 35./27 },
    { 1631./55296, 175./512, 575./13824, 44275./110592, 253./4096 }
  };

  template<class DerivativeFunctor, class Vector>
  const double RK5Iterator<DerivativeFunctor, Vector>::fC[] = {
    37./378, 250./621, 125./594, 0, 512./1771
  };

  template<class DerivativeFunctor, class Vector>
  const double RK5Iterator<DerivativeFunctor, Vector>::fD[] = {
    RK5Iterator<DerivativeFunctor, Vector>::fC[0] - 2825./27648,
    RK5Iterator<DerivativeFunctor, Vector>::fC[1] - 18575./48384,
    RK5Iterator<DerivativeFunctor, Vector>::fC[2] - 13525./55296,
    RK5Iterator<DerivativeFunctor, Vector>::fC[3] - 277./14336,
    RK5Iterator<DerivativeFunctor, Vector>::fC[4] - 0.25
  };


  class DefaultRK5ErrorScaling {
  public:
    template<class Vector>
    void
    operator()(const unsigned int n, const double dx,
               const Vector& y, const Vector& dYdX,
               Vector& errorScaling)
      const
    {
      for (unsigned int i = 0; i < n; ++i)
        errorScaling[i] = std::abs(y[i]) + std::abs(dx*dYdX[i]) + 1e-1;
    }
  };


  /**
    \class AdaptiveRK5Iterator RK5ODEIntegrator.h "utl/RK5ODEIntegrator.h"

    This iterator class will advance the solution from the current point
    maintaining the accuracy of the result.
    See class RK5ODEIntegrator for details.

    \author Darko Veberic
    \date 20 Jun 2008
    \version $Id$

    \ingroup math
  */
  template<class DerivativeFunctor,
           class Vector,
           class ErrorScalingPolicy = DefaultRK5ErrorScaling>
  class AdaptiveRK5Iterator :
    protected RK5Iterator<DerivativeFunctor, Vector>,
    public SafeBoolCast<AdaptiveRK5Iterator<DerivativeFunctor, Vector, ErrorScalingPolicy> > {
  public:
    AdaptiveRK5Iterator(const double x0, const double dx0,
                        const Vector& y0,
                        DerivativeFunctor& d,
                        const double accuracy = 1e-5,
                        const ErrorScalingPolicy& errorScaling =
                          DefaultRK5ErrorScaling()) :
      RK5Iterator<DerivativeFunctor, Vector>(x0, y0, d),
      fNextStepSize(dx0),
      fAccuracy(accuracy),
      fErrorScaling(errorScaling),
      fStatus(true)
    { }

    AdaptiveRK5Iterator(const double dx0,
                        const RK5Iterator<DerivativeFunctor, Vector>& rk5It,
                        const double accuracy = 1e-5,
                        const ErrorScalingPolicy& errorScaling =
                          DefaultRK5ErrorScaling()) :
      RK5Iterator<DerivativeFunctor, Vector>(rk5It),
      fNextStepSize(dx0),
      fAccuracy(accuracy),
      fErrorScaling(errorScaling),
      fStatus(true)
    { }

    double GetX() const { return this->fX; }

    Vector& GetY() { return this->fY; }

    const Vector& GetY() const { return this->fY; }

    const Vector& GetYError() const { return this->fYError; }

    bool BoolCast() const { return fStatus; }

    void ClearStatus() { fStatus = true; }

    bool
    Advance(const double dx)
    {
      double& x = this->fX;
      Vector& y = this->fY;
      Vector& yError = this->fYError;
      const unsigned int n = this->GetDerivativeFunctor();

      Vector dYdX;
      this->GetDerivativeFunctor()(x, y, dYdX);
      Vector y1;
      double dxAttempt = dx;
      double x1 = x + dxAttempt;
      Vector scaling;
      fErrorScaling(n, dx, y, dYdX, scaling);
      double relErr;
      while (true) {
        this->Step(dxAttempt, y1, dYdX);
        relErr =
          GetMaxScaledError(n, yError, scaling) / fAccuracy;
        if (relErr < 1)
          break;
        double stepFactor = fSlightlyLess * std::pow(relErr, fPowerDown);
        if (stepFactor < 0.1)
          stepFactor = 0.1;
        dxAttempt *= stepFactor;
        x1 = x + dxAttempt;
        if (x1 == x) {
          fNextStepSize = dx;
          fStatus = false;
          return false;
        }
      }
      fNextStepSize = dxAttempt *
        ((relErr > fPowerLawThreshold) ?
         fSlightlyLess * pow(relErr, fPowerUp) : 5);
      x = x1;
      this->Copy(n, y1, y);
      fStatus = true;
      return true;
    }

    bool
    Advance()
    {
      Advance(fNextStepSize);
      return fStatus;
    }

    AdaptiveRK5Iterator&
    operator+=(const double dx)
    { Advance(dx); return *this; }

    AdaptiveRK5Iterator&
    operator-=(const double dx)
    { Advance(-dx); return *this; }

    AdaptiveRK5Iterator&
    operator++()
    { Advance(); return *this; }

    AdaptiveRK5Iterator&
    operator++(int)
    { Advance(); return *this; }

  private:
    static
    double
    GetMaxScaledError(const unsigned int n,
                      const Vector& error, const Vector& scaling)
    {
      double max = std::abs(error[0] / scaling[0]);
      for (unsigned int i = 1; i < n; ++i) {
        const double se = std::abs(error[i] / scaling[i]);
        if (se > max)
          max = se;
      }
      return max;
    }

    double fNextStepSize;
    double fAccuracy;
    const ErrorScalingPolicy& fErrorScaling;
    bool fStatus;

    static const double fPowerUp;
    static const double fPowerDown;
    static const double fSlightlyLess;
    static const double fPowerLawThreshold;
  };

  template<
    class DerivativeFunctor,
    class Vector,
    class ErrorScalingPolicy
  > const double AdaptiveRK5Iterator<
    DerivativeFunctor,
    Vector,
    ErrorScalingPolicy
  >::fPowerUp = -0.2;

  template<
    class DerivativeFunctor,
    class Vector,
    class ErrorScalingPolicy
  > const double AdaptiveRK5Iterator<
    DerivativeFunctor,
    Vector,
    ErrorScalingPolicy
  >::fPowerDown = -0.25;

  template<
    class DerivativeFunctor,
    class Vector,
    class ErrorScalingPolicy
  > const double AdaptiveRK5Iterator<
    DerivativeFunctor,
    Vector,
    ErrorScalingPolicy
  >::fSlightlyLess = 0.9;

  template<
    class DerivativeFunctor,
    class Vector,
    class ErrorScalingPolicy
  > const double AdaptiveRK5Iterator<
    DerivativeFunctor,
    Vector,
    ErrorScalingPolicy
  >::fPowerLawThreshold = 1.889568e-4;


  /**
    \class RK5ODEIntegrator RK5ODEIntegrator.h "utl/RK5ODEIntegrator.h"

    This class prepares the integration of the Ordinary Differential Equations
    (ODE) of the form
    \f[
      \frac{{\rm d}y_i}{{\rm d}x} = f_i(x, y_0, \ldots, y_{n-1})
    \f]
    for \f$i\in{0, \ldots, n-1}\f$.

    ODE is defined through an function class with operator() defined
    (usualy such classes are called functor in C++). This functor has to
    define the operator() and unsigned int cast operator:
    \code
      class MyODE {
      public:
        void
        operator()(const double x, const double y[10], double dydx[10])
        {
          dydx[0] = ...
          dydx[1] = ...
        }

        operator unsigned int() { return 10; }
      };
    \endcode
    Operator() calculates derivatives \f${\rm d}y_i/{\rm d}x\f$ of variables
    \f$y_i\f$ with respect to the independent variable x.

    The cast operator unsigned int() returns number of OD equations (in this
    example 10).

    The class RK5ODEIntegrator returns an iterator that can be advanced with
    operator += dx for some step dx in independent variable.
    \code
      RK4ODEIntegrator<MyODE, double [10]> ode(MyODE());

      RK4Iterator<MyODE, double [10]> it = ode.Begin(x0, y0);

      it += dx;
    \endcode

    Next to the iterator (similar to the RK4Iterator) there is also
    AdaptiveRK5Iterator available that maintains the accuracy of the
    result by taking steps of optimal size;

    Note that the suitable candidates for template parameter Vector
    are vector classes with fixed size (std::vector<> is not appropriate
    since the default constructor leaves the object with zero size) that
    support the operator[] for element access. You can use the C-style
    arrays double[n] or Offline class utl::SVector<n>.

    See also utl::RK4ODEIntegrator and testODEIntegrator.cc for a real
    examples.

    \author Darko Veberic
    \date 20 Jun 2008
    \version $Id$

    \ingroup math
  */
  template<class DerivativeFunctor>
  class RK5ODEIntegrator {
  public:
    RK5ODEIntegrator(DerivativeFunctor& d)
      : fDerivativeFunctor(d) { }

    template<class Vector>
    RK5Iterator<DerivativeFunctor, Vector>
    Begin(const double x, const Vector& y)
    {
      return
        RK5Iterator
          <DerivativeFunctor, Vector>(x, y, fDerivativeFunctor);
    }

    template<class Vector>
    AdaptiveRK5Iterator<DerivativeFunctor, Vector>
    AdaptiveBegin(const double x, const double dx, const Vector& y,
                  const double accuracy = 1e-5)
    {
      return
        AdaptiveRK5Iterator<DerivativeFunctor, Vector>(
          x, dx, y, fDerivativeFunctor, accuracy
        );
    }

    template<class Vector,
             class ErrorScalingPolicy>
    AdaptiveRK5Iterator<DerivativeFunctor, Vector, ErrorScalingPolicy>
    AdaptiveBegin(const double x, const double dx, const Vector& y,
                  const double accuracy, ErrorScalingPolicy& es)
    {
      return
        AdaptiveRK5Iterator<DerivativeFunctor, Vector>(
          x, dx, y, fDerivativeFunctor, accuracy, es
        );
    }

    DerivativeFunctor& GetDerivativeFunctor()
    { return fDerivativeFunctor; }

  private:
    DerivativeFunctor& fDerivativeFunctor;
  };

}

#endif