# include <cppad/cppad.hpp>

// =========================================================================
// define types for each level
namespace { // BEGIN empty namespace
typedef CppAD::AD<double>     ADdouble;
typedef CppAD::AD< ADdouble > ADDdouble;

// -------------------------------------------------------------------------
// class definition for C++ function object that defines ODE
class Ode {
private:
	// copy of a that is set by constructor and used by g(y)
	CPPAD_TEST_VECTOR< ADdouble > x_; 
public:
	// constructor
	Ode( CPPAD_TEST_VECTOR< ADdouble > x) : x_(x)
	{ }
	// the function g(y) is evaluated with two levels of taping
	CPPAD_TEST_VECTOR< ADDdouble > operator()
	( const CPPAD_TEST_VECTOR< ADDdouble > &y) const
	{	size_t n = y.size();
		CPPAD_TEST_VECTOR< ADDdouble > g(n);
		size_t i;
		g[0] = x_[0];
		for(i = 1; i < n; i++)
			g[i] = x_[i] * y[i-1];

		return g;
	}
};

// -------------------------------------------------------------------------
// Routine that uses Taylor's method to solve ordinary differential equaitons
// and allows for algorithmic differentiation of the solution. 
CPPAD_TEST_VECTOR < ADdouble > taylor_ode(
	Ode                     G       ,  // function that defines the ODE
	size_t                  order   ,  // order of Taylor's method used
	size_t                  nstep   ,  // number of steps to take
	ADdouble                &dt     ,  // Delta t for each step
	CPPAD_TEST_VECTOR< ADdouble > &y_ini  )  // y(t) at the initial time
{
	// some temporary indices
	size_t i, k, ell;

	// number of variables in the ODE
	size_t n = y_ini.size();

	// copies of x and g(y) with two levels of taping
	CPPAD_TEST_VECTOR< ADDdouble >   Y(n), Z(n);

	// y, y^{(k)} , z^{(k)}, and y^{(k+1)}
	CPPAD_TEST_VECTOR< ADdouble >  y(n), y_k(n), z_k(n), y_kp(n);
	
	// initialize x
	for(i = 0; i < n; i++)
		y[i] = y_ini[i];

	// loop with respect to each step of Taylors method
	for(ell = 0; ell < nstep; ell++)
	{	// prepare to compute derivatives of in ADdouble
		for(i = 0; i < n; i++)
			Y[i] = y[i];
		CppAD::Independent(Y);

		// evaluate ODE in ADDdouble
		Z = G(Y);

		// define differentiable version of g: X -> Y
		// that computes its derivatives in ADdouble
		CppAD::ADFun<ADdouble> g(Y, Z);

		// Use Taylor's method to take a step
		y_k            = y;     // initialize y^{(k)}
		ADdouble dt_kp = dt;    // initialize dt^(k+1)
		for(k = 0; k <= order; k++)
		{	// evaluate k-th order Taylor coefficient of y
			z_k = g.Forward(k, y_k);
 
			for(i = 0; i < n; i++)
			{	// convert to (k+1)-Taylor coefficient for x
				y_kp[i] = z_k[i] / ADdouble(k + 1);

				// add term for to this Taylor coefficient
				// to solution for y(t, x)
				y[i]    += y_kp[i] * dt_kp;
			}
			// next power of t
			dt_kp *= dt;
			// next Taylor coefficient
			y_k   = y_kp;
		}
	}
	return y;
}
} // END empty namespace
// ==========================================================================
// Routine that tests alogirhtmic differentiation of solutions computed
// by the routine taylor_ode.
bool ode_taylor(void)
{	// initialize the return value as true	
	bool ok = true;

	// number of components in differential equation
	size_t n = 4;

	// some temporary indices
	size_t i, j;

	// parameter vector in both double and ADdouble
	CPPAD_TEST_VECTOR<double>   x(n);
	CPPAD_TEST_VECTOR<ADdouble> X(n);
	for(i = 0; i < n; i++)
		X[i] = x[i] = double(i + 1);

	// declare the parameters as the independent variable
	CppAD::Independent(X);

	// arguments to taylor_ode 
	Ode G(X);                // function that defines the ODE
	size_t   order = n;      // order of Taylor's method used
	size_t   nstep = 2;      // number of steps to take
	ADdouble DT    = 1.;     // Delta t for each step
	// value of y(t, x) at the initial time
	CPPAD_TEST_VECTOR< ADdouble > Y_INI(n);
	for(i = 0; i < n; i++)
		Y_INI[i] = 0.;

	// integrate the differential equation
	CPPAD_TEST_VECTOR< ADdouble > Y_FINAL(n);
 	Y_FINAL = taylor_ode(G, order, nstep, DT, Y_INI);

	// define differentiable fucntion object f : A -> Y_FINAL
	// that computes its derivatives in double
	CppAD::ADFun<double> f(X, Y_FINAL);

	// check function values
	double check = 1.;
	double t     = nstep * Value(DT);
	for(i = 0; i < n; i++)
	{	check *= x[i] * t / double(i + 1);
		ok &= CppAD::NearEqual(Value(Y_FINAL[i]), check, 1e-10, 1e-10);
	}

	// evaluate the Jacobian of h at a
	CPPAD_TEST_VECTOR<double> jac ( f.Jacobian(x) );
	// There appears to be a bug in g++ version 4.4.2 becasue it generates
	// a warning for the equivalent form
	// CPPAD_TEST_VECTOR<double> jac = f.Jacobian(x);

	// check Jacobian 
	for(i = 0; i < n; i++)
	{	for(j = 0; j < n; j++)
		{	double jac_ij = jac[i * n + j]; 
			if( i < j )
				check = 0.;
			else	check = Value( Y_FINAL[i] ) / x[j];
			ok &= CppAD::NearEqual(jac_ij, check, 1e-10, 1e-10);
		}
	}
	return ok;
}