#include <../../nrnconf.h>
// differential algebraic system solver interface to DDASPK

// DDASPK is translated from fortran with f2c. Hence all args are
// call by reference

#include <stdio.h>
#include <math.h>
#include <InterViews/resource.h>
extern "C" {
#include "spmatrix.h"
}
#include "nrnoc2iv.h"
#include "cvodeobj.h"
#include "nrndaspk.h"
#include "netcvode.h"
#include "ida/ida.h"
#include "ida/ida_impl.h"
#include "mymath.h"

// the state of the g - d2/dx2 matrix for voltages
#define INVALID 0
#define NO_CAP 1
#define SETUP 2
#define FACTORED 3
static int solve_state_;

// prototypes

double Daspk::dteps_;

extern "C" {
	
extern void linmod_dkres(double*, double*, double*);
extern void linmod_dkpsol(double);
extern void nrn_rhs(NrnThread*);
extern void nrn_lhs(NrnThread*);
extern void nrn_solve(NrnThread*);
void nrn_daspk_init_step(double, double, int);
// this is private in ida.c but we want to check if our initialization
// is good. Unfortunately ewt is set on the first call to solve which
// is too late for us.
extern booleantype IDAEwtSet(IDAMem IDA_mem, N_Vector ycur);

//extern double t, dt;
#define nt_dt nrn_threads->_dt
#define nt_t nrn_threads->_t

static void daspk_nrn_solve(NrnThread* nt) {
	nrn_solve(nt);
}

static int res_gvardt(
	realtype t, N_Vector y, N_Vector yp, N_Vector delta,
	void* rdata
);

static int minit(IDAMem);

static int msetup(IDAMem mem, N_Vector y, N_Vector ydot, N_Vector delta,
	N_Vector tempv1, N_Vector tempv2, N_Vector tempv3
);

static int msolve(IDAMem mem, N_Vector b, N_Vector ycur, N_Vector ypcur,
	N_Vector deltacur
);

static int mfree(IDAMem);

}


// at least in DARWIN the following is already declared so avoid conflict
#define thread_t nrn_thread_t

// residual
static N_Vector nvec_y;
static N_Vector nvec_yp;
static N_Vector nvec_delta;
static double thread_t;
static double thread_cj;
static int thread_ier;
static Cvode* thread_cv;
static void* res_thread(NrnThread* nt) {
	int i = nt->id;
	Cvode* cv = thread_cv;
	int ier = cv->res(thread_t, cv->n_vector_data(nvec_y, i),
		cv->n_vector_data(nvec_yp, i), cv->n_vector_data(nvec_delta, i), nt);
	if (ier != 0) {
		thread_ier = ier;
	}
	return 0;
}
static int res_gvardt(
	realtype t, N_Vector y, N_Vector yp, N_Vector delta,
	void* rdata
){
	thread_cv = (Cvode*)rdata;
	nvec_y = y;
	nvec_yp = yp;
	nvec_delta = delta;
	thread_t = t;
	thread_ier = 0;
	nrn_multithread_job(res_thread);
	return thread_ier;
}

// linear solver specific allocation and initialization
static int minit(IDAMem) {
	return IDA_SUCCESS;
}

// linear solver preparation for subsequent calls to msolve
// approximation to jacobian. Everything necessary for solving P*x = b
static int msetup(IDAMem mem, N_Vector y, N_Vector yp, N_Vector,
	N_Vector, N_Vector, N_Vector
){
	Cvode* cv = (Cvode*)mem->ida_rdata;
	++cv->jac_calls_;
	return 0;
}

/* solve P*x = b */
static void* msolve_thread(NrnThread* nt) {
	int i = nt->id;
	Cvode* cv = thread_cv;
	int ier = cv->psol(thread_t, cv->n_vector_data(nvec_y, i),
		cv->n_vector_data(nvec_yp, i), thread_cj, nt);
	if (ier != 0) {
		thread_ier = ier;
	}
	return 0;
}
static int msolve(IDAMem mem, N_Vector b, N_Vector w, N_Vector ycur, N_Vector, N_Vector){
	thread_cv = (Cvode*)mem->ida_rdata;
	thread_t = mem->ida_tn;
	nvec_y = ycur;
	nvec_yp = b;
	thread_cj = mem->ida_cj;
	nrn_multithread_job(msolve_thread);
	return thread_ier;
}

static int mfree(IDAMem) {return IDA_SUCCESS;}

Daspk::Daspk(Cvode* cv, int neq) {
//	printf("Daspk::Daspk\n");
	cv_ = cv;
	yp_ = cv->nvnew(neq);
	delta_ = cv->nvnew(neq);
	parasite_ = cv->nvnew(neq);
	use_parasite_ = false;
	spmat_ = nil;
	mem_ = nil;
}

Daspk::~Daspk() {
//	printf("Daspk::~Daspk\n");
	N_VDestroy(delta_);
	N_VDestroy(yp_);
	if (mem_) {
		IDAFree((IDAMem)mem_);
	}
}

void Daspk::ida_init() {
	int ier;
	if (mem_) {
		ier = IDAReInit(mem_, res_gvardt, cv_->t_, cv_->y_, yp_,
			IDA_SV, &cv_->ncv_->rtol_, cv_->atolnvec_
		);
		if (ier < 0) {
			hoc_execerror("IDAReInit error", 0);
		}
	}else{
		IDAMem mem = (IDAMem) IDACreate();
		if (!mem) {
			hoc_execerror("IDAMalloc error", 0);
		}
		IDASetRdata(mem, cv_);
		ier = IDAMalloc(mem, res_gvardt, cv_->t_, cv_->y_, yp_,
			IDA_SV, &cv_->ncv_->rtol_, cv_->atolnvec_
		);
		mem->ida_linit = minit;
		mem->ida_lsetup = msetup;
		mem->ida_lsolve = msolve;
		mem->ida_lfree = mfree;
		mem->ida_setupNonNull = false;
		mem_ = mem;
	}
}

void Daspk::info() {
}


// last two bits, 0 error, 1 warning, 2 apply parasitic
// if init_failure_style & 010, then use the original method
int Daspk::init_failure_style_;
int Daspk::init_try_again_;
int Daspk::first_try_init_failures_;

static void* do_ode_thread(NrnThread* nt) {
	int i;
	Cvode* cv = thread_cv;
	cv->do_ode(nt);
	CvodeThreadData& z = cv->ctd_[nt->id];
	double* yp = cv->n_vector_data(nvec_yp, nt->id);
	for (i=z.neq_v_; i < z.nvsize_; ++i) {
		yp[i] = *(z.pvdot_[i]);
	}
	return 0;
}

int Daspk::init() {
	int i;
#if 0
printf("Daspk_init t_=%20.12g t-t_=%g t0_-t_=%g\n",
cv_->t_, t-cv_->t_, cv_->t0_-cv_->t_);
#endif
	N_VConst(0., yp_);

	// the new initial condition is based on a dteps_ step backward euler
	// linear solution with respect to the old state in order to
	// start the following initial condition calculation with a "valid"
	// (in a linear system sense) initial state.

	double tt = cv_->t_;
	double dtinv = 1./dteps_;
    if (init_failure_style_ & 010) {
	cv_->play_continuous(tt);
	nrn_daspk_init_step(tt, dteps_, 1);
	nrn_daspk_init_step(tt, dteps_, 1);
	cv_->daspk_gather_y(yp_);
	cv_->play_continuous(tt);
	nrn_daspk_init_step(tt, dteps_, 1);
	cv_->daspk_gather_y(cv_->y_);
	N_VLinearSum(dtinv, cv_->y_, -dtinv, yp_, yp_);
    }else{
#if 0
	cv_->play_continuous(tt);
	nrn_daspk_init_step(tt, dteps_, 1);
	cv_->daspk_gather_y(cv_->y_);
	tt = cv_->t_ + dteps_;
	cv_->play_continuous(tt);
	nrn_daspk_init_step(tt, dteps_, 1);
	cv_->daspk_gather_y(yp_);
	N_VLinearSum(dtinv, yp_, -dtinv, cv_->y_, yp_);
	cv_->daspk_scatter_y(cv_->y_);
#else
	cv_->play_continuous(tt);
	nrn_daspk_init_step(tt, dteps_, 1); // first approx to y (and maybe good enough)
	nrn_daspk_init_step(tt, dteps_, 1); // 2nd approx to y (trouble with 2sramp.hoc)

	cv_->daspk_gather_y(cv_->y_);
	tt = cv_->t_ + dteps_;
	cv_->play_continuous(tt);
	nrn_daspk_init_step(tt, dteps_, 0); // rhs contains delta y (for v, vext, linmod
	cv_->gather_ydot(yp_);
	N_VScale(dtinv, yp_, yp_);
#endif
    }
	thread_cv = cv_;
	nvec_yp = yp_;
	nrn_multithread_job(do_ode_thread);
	ida_init();
#if 1
	// test
//printf("test\n");
	if (!IDAEwtSet((IDAMem)mem_, cv_->y_)) {
		hoc_execerror("Bad Ida error weight vector", 0);
	}
	use_parasite_ = false;
	res_gvardt(tt, cv_->y_, yp_, parasite_, cv_);
	double norm = N_VWrmsNorm(parasite_, ((IDAMem)mem_)->ida_ewt);
//printf("norm=%g at t=%g\n", norm, t);
	if (norm > 1.) {
	    switch (init_failure_style_ & 03) {
	    case 0:
		printf("IDA initialization failure, weighted norm of residual=%g\n", norm);
		return IDA_ERR_FAIL;
		break;
	    case 1:
		printf("IDA initialization warning, weighted norm of residual=%g\n", norm);
		break;
	    case 2:
		printf("IDA initialization warning, weighted norm of residual=%g\n", norm);
		use_parasite_ = true;
		t_parasite_ = nt_t;
		printf("  subtracting (for next 1e-6 ms): f(y', y, %g)*exp(-1e7*(t-%g))\n", nt_t, nt_t);
		break;
	    }
#if 0
for (i=0; i < cv_->neq_; ++i) {
	printf("   %d %g %g %g %g\n", i, nt_t, N_VGetArrayPointer(cv_->y_)[i], N_VGetArrayPointer(yp_)[i], N_VGetArrayPointer(delta_)[i]);
}
#endif
		if (init_try_again_ < 0) {
			++first_try_init_failures_;
			init_try_again_ += 1;
			int err = init();
			init_try_again_ = 0;
			return err;
		}
return 0;
	}
#endif

	return 0;
}

int Daspk::advance_tn(double tstop) {
//printf("Daspk:solve %s tin=%g tstop=%g\n", (info_[0] == 0)?"initial":"continue", *pt_, rwork_[0]);
	double tn = cv_->tn_;
	IDASetStopTime(mem_, tstop);
	int ier = IDASolve(mem_, tstop, &cv_->t_, cv_->y_, yp_, IDA_ONE_STEP_TSTOP);
	if (ier < 0) {
		//printf("DASPK advance_tn error\n");
		return ier;
	}
#if 0
	if (ier > 0 && t < cv_->t_) {
		// interpolation to tstop does not call res. So we have to.
		cv_->res(cv_->t_, N_VGetArrayPointer(cv_->y_), N_VGetArrayPointer(yp_), N_VGetArrayPointer(delta_));
		assert(MyMath::eq(t, cv_->t_, NetCvode::eps(t)));
	}
#else
	// this is very bad, performance-wise. However ida modifies its states
	// after a call to fun with the proper t.
	res_gvardt(cv_->t_, cv_->y_, yp_, delta_, cv_);
#endif
	cv_->t0_ = tn;
	cv_->tn_ = cv_->t_;
//printf("Daspk:solve %s tout=%g tstop-tout=%g idid=%d\n", (info_[0] == 0)?"initial":"continue", *pt_, rwork_[0] - *pt_, idid_);
	return ier;
}

int Daspk::interpolate(double tt) {
//printf("Daspk::interpolate %.15g\n", tt);
	assert (tt >= cv_->t0_ && tt <= cv_->tn_);
	IDASetStopTime(mem_, tt);
	int ier = IDASolve(mem_, tt, &cv_->t_, cv_->y_, yp_, IDA_NORMAL);
	if (ier < 0) {
		printf("DASPK interpolate error\n");
		return ier;
	}
	assert(MyMath::eq(tt, cv_->t_, NetCvode::eps(cv_->t_)));
	// interpolation does not call res. So we have to.
	res_gvardt(cv_->t_, cv_->y_, yp_, delta_, cv_);
//if(MyMath::eq(t, cv_->t_, NetCvode::eps(cv_->t_))) {
//printf("t=%.15g t_=%.15g\n", t, cv_->t_);
//}
//	assert(MyMath::eq(t, cv_->t_, NetCvode::eps(cv_->t_)));
	return ier;
}

void Daspk::statistics() {
#if 0
	printf("rwork size = %d\n", iwork_[18-1]);
	printf("iwork size = %d\n", iwork_[17-1]);
	printf("Number of time steps = %d\n", iwork_[11-1]);
	printf("Number of residual evaluations = %d\n", iwork_[12-1]);
	printf("Number of Jac evaluations = %d\n", iwork_[13-1]);
	printf("Number of preconditioner solves = %d\n", iwork_[21-1]);
	printf("Number of nonlinear iterations = %d\n", iwork_[19-1]);
	printf("Number of linear iterations = %d\n", iwork_[20-1]);
	double avlin = double(iwork_[20-1])/double(iwork_[19-1]);
	printf("Average Krylov subspace dimension = %g\n", avlin);
	printf("nonlinear conv. failures = %d\n", iwork_[15-1]);
	printf("linear conv. failures = %d\n", iwork_[16-1]);
#endif
	if (first_try_init_failures_) {
		printf("   %d First try Initialization failures\n", first_try_init_failures_);
	}
}

static void* daspk_scatter_thread(NrnThread* nt) {
	thread_cv->daspk_scatter_y(thread_cv->n_vector_data(nvec_y, nt->id), nt->id);
	return 0;
}
void Cvode::daspk_scatter_y(N_Vector y) {
	thread_cv = this;
	nvec_y = y;
	nrn_multithread_job(daspk_scatter_thread);
}
void Cvode::daspk_scatter_y(double* y, int tid) {
	// the dependent variables in daspk are vi,vx,etc
	// whereas in the node structure we need vm, vx
	// note that a corresponding transformation for gather_ydot is
	// not needed since the matrix solve is already with respect to vi,vx
	// in all cases. (i.e. the solution vector is in the right hand side
	// and refers to vi, vx.
	scatter_y(y, tid);
	// transform the vm+vext to vm
	CvodeThreadData& z = ctd_[tid];
	if (z.cmlext_) {
		Memb_list* ml = z.cmlext_->ml;
		int i, n = ml->nodecount;
		for (i=0; i < n; ++i) {
			Node* nd = ml->nodelist[i];
			NODEV(nd) -= nd->extnode->v[0];
		}
	}
}
static void* daspk_gather_thread(NrnThread* nt) {
	thread_cv->daspk_gather_y(thread_cv->n_vector_data(nvec_y, nt->id), nt->id);
	return 0;
}
void Cvode::daspk_gather_y(N_Vector y) {
	thread_cv = this;
	nvec_y = y;
	nrn_multithread_job(daspk_gather_thread);
}
void Cvode::daspk_gather_y(double* y, int tid) {
	gather_y(y, tid);
	// transform vm to vm+vext
	CvodeThreadData& z = ctd_[tid];
	if (z.cmlext_) {
		Memb_list* ml = z.cmlext_->ml;
		int i, n = ml->nodecount;
		for (i=0; i < n; ++i) {
			Node* nd = ml->nodelist[i];
			int j = nd->eqn_index_;
			y[j-1] += y[j];
		}
	}
}

// for res and psol the equations for c*yp = f(y) are
// cast in the form G(t,y,yp) = f(y) - c*yp
// So res calculates delta = f(y) - c*yp
// and psol solves (c*cj - df/dy)*x = -b
// Note that since cvode uses  J = 1 - gam*df/dy and
// ida uses J = df/dy - cj*df/dyp that is the origin of the use of -b in our
// psol and also why all the non-voltage odes are scaled by dt at the
// end of it.

int Cvode::res(double tt, double* y, double* yprime, double* delta, NrnThread* nt) {
	CvodeThreadData& z = ctd_[nt->id];
	++f_calls_;
	static int res_;
	int i;
	nt->_t = tt;
	res_++;

#if 0
printf("Cvode::res enter tt=%g\n", tt);
for (i=0; i < z.nvsize_; ++i) {
	printf("   %d %g %g %g\n", i, y[i], yprime[i], delta[i]);
}
#endif
	nt->_vcv = this; // some models may need to know this
	daspk_scatter_y(y, nt->id); // vi, vext, channel states, linmod non-node y.
	// rhs of cy' = f(y)
	play_continuous_thread(tt, nt);
	nrn_rhs(nt);
	do_ode(nt);
	// accumulate into delta
	gather_ydot(delta, nt->id);

	// now calculate -c*yp. i.e.
	// cm*vm' + c_linmod*vi' internal current balance
	// cx*vx' + c_linmod*vx' external current balance
	// c_linmod*y' non-node linmod states
	// y' mechanism states

	// this can be accumulated into delta in several stages
	// -cm*vm'and -cx*vx for current balance equation delta's
	// -c_linmod*yp (but note that the node yp yp(vm)+yp(vx))
	// subtract yp from mechanism state delta's

#if 0
printf("Cvode::res after ode and gather_ydot into delta\n");
for (i=0; i < z.nvsize_; ++i) {
	printf("   %d %g %g %g\n", i, y[i], yprime[i], delta[i]);
}
#endif
	// the cap nodes : see nrnoc/capac.c for location of cm, etc.
	// these are not in same order as for cvode but are in
	// spmatrix order mixed with nocap nodes and extracellular
	// therefore we use the Node.eqn_index to calculate the delta index.
//	assert(use_sparse13 == true && nlayer <= 1);
	assert(use_sparse13 == true);
	if (z.cmlcap_) {
		Memb_list* ml = z.cmlcap_->ml;
		int n = ml->nodecount;
		for (i=0; i < n; ++i) {
			double* cd = ml->data[i];
			Node* nd = ml->nodelist[i];
			int j = nd->eqn_index_ - 1;
			Extnode* nde = nd->extnode;
			if (nde) {
				double cdvm = 1e-3 * cd[0] * (yprime[j] - yprime[j+1]);
				delta[j] -= cdvm;
				delta[j+1] += cdvm;
				// i_cap
				cd[1] = cdvm;
#if I_MEMBRANE
				// add i_cap to i_ion which is in sav_g
				// this will be copied to i_membrane below
				nde->param[3+3*nlayer] += cdvm;
#endif
			}else{
				double cdvm = 1e-3 * cd[0] * yprime[j];
				delta[j] -= cdvm;
				cd[1] = cdvm;
			}
		}
	}
	// See nrnoc/excelln.c for location of cx.
	if (z.cmlext_) {
		Memb_list* ml = z.cmlext_->ml;
		int n = ml->nodecount;
		for (i=0; i < n; ++i) {
			double* cd = ml->data[i];
			Node* nd = ml->nodelist[i];
			int j = nd->eqn_index_;
#if I_MEMBRANE
			// i_membrane = sav_rhs --- even for zero area nodes
			cd[1+3*nlayer] = cd[3+3*nlayer];
#endif
#if EXTRACELLULAR == 1
			// only works for one layer
			// otherwise loop over layer,
			// xc is (pd + 2*(nlayer))[layer]
			// and deal with yprime[i+layer]-yprime[i+layer+1]
			delta[j] -= 1e-3 * cd[2] * yprime[j];
#else
			int k, jj;
			double x;
			k = nlayer-1;
			jj = j+k;
			delta[jj] -= 1e-3*cd[2*nlayer+k]*(yprime[jj]);
			for (k=nlayer-2; k >= 0; --k) {
				// k=0 refers to stuff between layer 0 and 1
				// j is for vext[0]				
				jj = j+k;
				x = 1e-3*cd[2*nlayer+k]*(yprime[jj] - yprime[jj+1]);
				delta[jj] -= x;
				delta[jj+1] += x; // last one in iteration is nlayer-1
			}
			
#endif
		}
	}

	linmod_dkres(y, yprime, delta);

	// the ode's
	for (i=z.neq_v_; i < z.nvsize_; ++i) {
		delta[i] -= yprime[i];
	}

	for (i=0; i < z.nvsize_; ++i) {
		delta[i] *= -1.;
	}
	if (daspk_->use_parasite_ && tt - daspk_->t_parasite_ < 1e-6) {
		double fac = exp(1e7*(daspk_->t_parasite_ - tt));
		double* tps = n_vector_data(daspk_->parasite_, nt->id);
		for (i=0; i < z.nvsize_; ++i) {
			delta[i] -= tps[i]*fac;
		}
	}
	before_after(z.after_solve_, nt);
#if 0
printf("Cvode::res exit res_=%d tt=%20.12g\n", res_, tt);
for (i=0; i < z.nvsize_; ++i) {
	printf("   %d %g %g %g\n", i, y[i], yprime[i], delta[i]);
}
#endif
	nt->_vcv = 0;
#if 0
double e = 0;
for (i=0; i < z.nvsize_; ++i) {
	e += delta[i]*delta[i];
}
printf("Cvode::res %d e=%g t=%.15g\n", res_, e, tt);
#endif
	return 0;
}

int Cvode::psol(double tt, double* y, double* b, double cj, NrnThread* _nt) {
	CvodeThreadData& z = ctd_[_nt->id];
	++mxb_calls_;
	int i;
	_nt->_t = tt;

#if 0
printf("Cvode::psol tt=%g solvestate=%d \n", tt, solve_state_);
for (i=0; i < z.nvsize_; ++i) {
printf(" %g", b[i]);
}
printf("\n");
#endif

	_nt->cj = cj;
	_nt->_dt = 1./cj;

	_nt->_vcv = this;
	daspk_scatter_y(y, _nt->id); // I'm not sure this is necessary.
	if (solve_state_ == INVALID) {
		nrn_lhs(_nt); // designed to setup M*[dvm+dvext, dvext, dy] = ...
		solve_state_ = SETUP;
	}
	if (solve_state_ == SETUP) {
		// if using sparse 13 then should factor
		solve_state_ = FACTORED;
	}
	scatter_ydot(b, _nt->id);
#if 0
printf("before nrn_solve matrix cj=%g\n", cj);
spPrint((char*)sp13mat_, 1,1,1);
printf("before nrn_solve actual_rhs=\n");
for (i=0; i < z.neq_v_; ++i) {
	printf("%d %g\n", i+1, actual_rhs[i+1]);
}
#endif
	daspk_nrn_solve(_nt); // not the cvode one
#if 0
//printf("after nrn_solve matrix\n");
//spPrint((char*)sp13mat_, 1,1,1);
printf("after nrn_solve actual_rhs=\n");
for (i=0; i < neq_v_; ++i) {
	printf("%d %g\n", i+1, actual_rhs[i+1]);
}
#endif
	solve_state_ = INVALID; // but not if using sparse13
	solvemem(_nt);
	gather_ydot(b, _nt->id);
	// the ode's of the form m' = (minf - m)/mtau in model descriptions compute
	// b = b/(1 + dt*mtau) since cvode required J = 1 - gam*df/dy
	// so we need to scale those results by 1/cj.
	for (i=z.neq_v_; i < z.nvsize_; ++i) {
		b[i] *= _nt->_dt;
	}
#if 0
for (i=0; i < z.nvsize_; ++i) {
printf(" %g", b[i]);
}
printf("\n");
#endif
	_nt->_vcv = 0;
	return 0;
}

N_Vector Daspk::ewtvec() {
	return ((IDAMem)mem_)->ida_ewt;
}

N_Vector Daspk::acorvec() {
	return ((IDAMem)mem_)->ida_delta;
}