# include <adolc/adouble.h>
# include <adolc/taping.h>
# include <adolc/drivers/drivers.h>

// definitions not in Adolc distribution and required to use CppAD::AD<adouble>
# include "base_adolc.hpp"

# include <cppad/cppad.hpp>
// ==========================================================================
namespace { // BEGIN empty namespace
// define types for each level
typedef adouble            ADdouble;
typedef CppAD::AD<adouble> 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_adolc(
	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 algorithmic differentiation of solutions computed
// by the routine taylor_ode.
bool ode_taylor_adolc(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;

	// set up for omp_alloc memory allocator
	using CppAD::omp_alloc; // the allocator
	size_t capacity;        // capacity of an allocation

	// the vector x with lenght n (or greater) in double 
	double* x = omp_alloc::create_array<double>(n, capacity);
	// the vector x with lenght n in ADouble
	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
	int tag = 0;                     // Adolc setup
	int keep = 1;
	trace_on(tag, keep);
	for(i = 0; i < n; i++)
		X[i] <<= double(i + 1);  // X is independent for adouble type

	// arguments to taylor_ode_adolc 
	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_adolc(G, order, nstep, DT, Y_INI);

	// declare the differentiable fucntion f : A -> Y_FINAL
	// (corresponding to the tape of adouble operations)
	double* y_final = omp_alloc::create_array<double>(n, capacity);
	for(i = 0; i < n; i++)
		Y_FINAL[i] >>= y_final[i];
	trace_off();

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

	// memory where Jacobian will be returned
	double* jac_ = omp_alloc::create_array<double>(n * n, capacity); 
	double** jac = omp_alloc::create_array<double*>(n, capacity);
	for(i = 0; i < n; i++)
		jac[i] = jac_ + i * n;

	// evaluate Jacobian of h at a
	size_t m = n;              // # dependent variables
	jacobian(tag, int(m), int(n), x, jac); 
	
	// check Jacobian 
	for(i = 0; i < n; i++)
	{	for(j = 0; j < n; j++)
		{	if( i < j )
				check = 0.;
			else	check = y_final[i] / x[j];
			ok &= CppAD::NearEqual(jac[i][j], check, 1e-10, 1e-10);
		}
	}

	// make memroy avaiable for other use by this thread
	omp_alloc::delete_array(x);
	omp_alloc::delete_array(y_final);
	omp_alloc::delete_array(jac_);
	omp_alloc::delete_array(jac);
	return ok;
}