#include <../../nrnconf.h>
#include <cstdio>
#include "nrndae.h"
#include "nrnoc2iv.h"

extern "C" {
	extern void nrndae_alloc();
	extern int nrndae_extra_eqn_count();
	extern void nrndae_init();
	extern void nrndae_rhs(); // relative to c*dy/dt = -g*y + b
	extern void nrndae_lhs();
	extern void nrndae_dkmap(double**, double**);
	extern void nrndae_dkres(double*, double*, double*);
	extern void nrndae_dkpsol(double);
	extern void nrndae_update();
	extern void nrn_matrix_node_free();
	extern int cvode_active_;
	extern int nrn_use_daspk_;
	extern int secondorder;
	extern int nrndae_list_is_empty();
}

static NrnDAEPtrList nrndae_list;

int nrndae_list_is_empty() {
	return nrndae_list.empty() ? 1 : 0;
}


void nrndae_register(NrnDAE* n) {
	nrndae_list.push_back(n);
}

void nrndae_deregister(NrnDAE* n) {
	nrndae_list.remove(n);
}

int nrndae_extra_eqn_count() {
	int neqn = 0;
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		neqn += (*m) -> extra_eqn_count();
	}
	return neqn;
}

void nrndae_update() {
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		(*m) -> update();		
	}
}

void nrndae_alloc() {
	NrnThread* _nt = nrn_threads;
	nrn_thread_error("NrnDAE only one thread allowed");
	int neqn = _nt -> end;
	if (_nt -> _ecell_memb_list) {
		neqn += _nt -> _ecell_memb_list -> nodecount * nlayer;
	}
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		(*m) -> alloc(neqn + 1);
		neqn += (*m) -> extra_eqn_count();
	}
}


void nrndae_init() {
	if ((!nrndae_list.empty()) &&
(secondorder > 0 || ((cvode_active_ > 0) && (nrn_use_daspk_ == 0)))) {
hoc_execerror("NrnDAEs only work with secondorder==0 or daspk",0);
	}
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		(*m) -> init();		
	}
}

void nrndae_rhs() {
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		(*m) -> rhs();
	}
}

void nrndae_lhs() {
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		(*m) -> lhs();
	}
}

void nrndae_dkmap(double** pv, double** pvdot) {
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		(*m) -> dkmap(pv, pvdot);
	}
}

void nrndae_dkres(double* y, double* yprime, double* delta) {
	// c*y' = f(y) so
	// delta = c*y' - f(y)
	for (NrnDAEPtrListIterator m = nrndae_list.begin(); m != nrndae_list.end(); m++) {
		(*m) -> dkres(y, yprime, delta);
	}
}

inline void NrnDAE::alloc_(int size, int start, int nnode, Node** nodes, int* elayer) {
}

void NrnDAE::alloc(int start_index) {
//printf("NrnDAE::alloc %lx\n", (long)this);
	size_ = y_.capacity();
	if (y0_) {
		assert(y0_ -> capacity() == size_);
	}
	assert(c_->nrow() == size_ && c_->ncol() == size_);
	cyp_.resize(size_);
	yptmp_.resize(size_);
	start_ = start_index;
//printf("start=%d size=%d\n", start_, size_);
	delete [] bmap_;
	bmap_ = new int[size_];
	for (int i = 0; i < size_; ++i) {
		if (i < nnode_) {
			bmap_[i] = nodes_[i]->eqn_index_ + elayer_[i];
			if (elayer_[i] > 0 && !nodes_[i]->extnode) {
//hoc_execerror(secname(nodes_[i]->sec), "NrnDAE: Referring to an extracellular layer but\nextracellular is not inserted.");
// instead treat as though connected to ground.
				bmap_[i] = 0;
			}
		}else{
			bmap_[i] = start_ + i - nnode_;
		}
	}
//printf("c_->alloc start=%d, nnode=%d\n", start_, nnode_);
	c_->alloc(start_, nnode_, nodes_, elayer_);
	
	// allow subclasses to do their own allocations as well
	alloc_(size_, start_, nnode_, nodes_, elayer_);
}

NrnDAE::NrnDAE(Matrix* cmat, Vect* const yvec, Vect* const y0, int nnode,
                   Node** const nodes, Vect* const elayer, void (*f_init)(void* data), void* const data)
        : y_(*yvec), yptmp_((Object*) NULL), cyp_((Object*) NULL), f_init_(f_init), data_(data) {
//printf("NrnDAE %lx\n", (long)this);
    if (cmat) {
        assumed_identity_ = NULL;
    } else {
        const int size = y_.capacity();
        assumed_identity_ = new OcSparseMatrix(size, size);
		//assumed_identity_->setdiag(0, 1);
        for (int i = 0; i < size; i++) (*assumed_identity_)(i, i) = 1;
        cmat = assumed_identity_;
    }
    c_ = new MatrixMap(cmat);
    Vect &elay = *elayer;
	nnode_ = nnode;
	nodes_ = nodes;
	if (nnode_ > 0) {
		elayer_ = new int[nnode_];
		if (elayer) {
			for (int i = 0; i < nnode_; ++i) {
				elayer_[i] = int(elay[i]);
			}
		} else {
			for (int i = 0; i < nnode_; ++i) {
				elayer_[i] = 0;
			}
		}
	} else {
		elayer_ = NULL;
	}
	y0_ = y0;
	bmap_ = new int[1];
	
	nrndae_register(this);
//	use_sparse13 = 1;
	nrn_matrix_node_free();
}


NrnDAE::~NrnDAE() {
	nrndae_deregister(this);
	delete [] bmap_;
    delete c_;
    delete assumed_identity_;
	if (elayer_) {
		delete [] elayer_;
	}
//	if (nrndae_list->count() == 0) {
//		use_sparse13 = 0;
//	}
	nrn_matrix_node_free();
}


int NrnDAE::extra_eqn_count() {
//printf("NrnDAE::extra_eqn_count %lx\n", (long)this);
//printf("  nnode_=%d g_->nrow()=%d\n", nnode_, g_->nrow());
	return c_->nrow() - nnode_;
}

void NrnDAE::dkmap(double** pv, double** pvdot) {
    //printf("NrnDAE::dkmap\n");
	NrnThread* _nt = nrn_threads;
	for (int i = nnode_; i < size_; ++i) {
//printf("bmap_[%d] = %d\n", i, bmap_[i]);
		pv[bmap_[i] - 1] = y_.vec() + i;
		pvdot[bmap_[i] - 1] = _nt->_actual_rhs + bmap_[i];
	}
}

void NrnDAE::update() {
//printf("NrnDAE::update %lx\n", (long)this);
	NrnThread* _nt = nrn_threads;
	// note that the following is correct also for states that refer
	// to the internal potential of a segment. i.e rhs is v + vext[0]
	for (int i = 0; i < size_; ++i) {
		y_[i] += _nt->_actual_rhs[bmap_[i]];
	}
//for (int i=0; i < size_; ++i) printf(" i=%d bmap_[i]=%d y_[i]=%g\n", i, bmap_[i], y_->elem(i));
}

void NrnDAE::init() {
//printf("NrnDAE::init %lx\n", (long)this);
//printf("init size_=%d %d %d %d\n", size_, y_->capacity(), y0_->capacity(), b_->capacity());
	Vect& y0 = *y0_;
	
	v2y();
    if (f_init_) {
        f_init_(data_);
    } else {
        if (y0_) {
            for (int i = nnode_; i < size_; ++i) {
                y_[i] = y0[i];
            }
        } else {
            for (int i = nnode_; i < size_; ++i) {
                y_[i] = 0.;
            }
        }
    }
//for (i=0; i < nnode_; ++i) printf(" i=%d y[i]=%g\n", i, y[i]);
}

void NrnDAE::v2y() {
	// vm,vext may be reinitialized between fixed steps and certainly
	// have been adjusted by daspk
	
	for (int i = 0; i < nnode_; ++i) {
		Node* nd = nodes_[i];
		//Note elayer_[0] refers to internal.
		if (elayer_[i] == 0) {
			y_[i] = NODEV(nd);
			if (nd -> extnode) {
				y_[i] += nd->extnode->v[0];
			}
		} else if (nd -> extnode) {
			y_[i] = nd -> extnode -> v[elayer_[i] - 1];
		}
	}
}



void NrnDAE::dkres(double* y, double* yprime, double* delta) {
//printf("NrnDAE::dkres %lx\n", (long)this);
	// delta is already f(y)
	// now subtract c*y'
	// The problem is the map between y and yprime and the local
	// representation of the y and yprime
	
	for (int i = 0; i < size_; ++i) {
//printf("%d %d %g %g\n", i, bmap_[i]-1, y[bmap_[i]-1], yprime[bmap_[i]-1]);
		yptmp_[i] = yprime[bmap_[i] - 1];
	}
    Vect* cyp;
    if (assumed_identity_) {
        // if c_ is KNOWN to be the identity matrix, then no multiplication to do
        // for now, this happens only when cmat = NULL
        // note that the user might change cmat, so we can't assume that if it's initially
        // the identity that it will stay that way
        cyp = &yptmp_;
    } else {
        c_->mulv(yptmp_, cyp_); // mulv cannot multiply in place
        cyp = &cyp_;
    }
	for (int i = 0; i < size_; ++i) {
		delta[bmap_[i] - 1] -= (*cyp)[i];
	}
}


void NrnDAE::rhs() {
	NrnThread* _nt = nrn_threads;
	v2y();
	f_(y_, yptmp_, size_);
	for (int i = 0; i < size_; ++i) {
		_nt->_actual_rhs[bmap_[i]] += yptmp_[i];
	}
}

void NrnDAE::lhs() {
//printf("NrnDAE::lhs %lx\n", (long)this);
//printf("  nrn_threads[0].cj = %g\n", nrn_threads[0].cj);
	//left side portion of (c/dt - J)*[dy] =  f(y)
	c_->add(nrn_threads[0].cj);
	v2y();
	jacobian_(y_)->add(jacobian_multiplier_() * -1);
}


