library(fields)
library(deSolve)

#################### functions for drawing phase planes and trajectories ###############

# draw phase plane
# input:
#	ode.sys:  ode system
#	ode.event, ode.root: event and root function for deSolve in the case that ode.sys is discountinous
# 	nullcline: function to compute nullcline
#	params: params for ode system
#	xrange: c(xmin,xmax) of the graph
#	yrange: c(ymin,ymax) of the graph
#	xstep(ystep): the x(y) distance between two vector tails
#	x.ini, y.ini: initial conditions of the system
#	times: time sequence of trajectories
draw.phase.plane <- function(ode.sys, ode.event, ode.root, 
					nullcline, params, 
					xrange, yrange, xstep, ystep,
					xlabel, ylabel, 
					x.ini, y.ini, times) {

	####### draw vector field #######
	# compute the tails of vectors
	x <- seq(xrange[1], xrange[2], xstep)
	y <- seq(yrange[1], yrange[2], ystep)
	vec.tail.x <- rep(x, time = length(y))
	vec.tail.y <- rep(y, each = length(x))

	# compute the directions of vectors
	n.vecs <- length(vec.tail.x)
	vec.dir.x <- rep(0, n.vecs)
	vec.dir.y <- rep(0, n.vecs)

	for (i in 1:n.vecs) {
		dir <- ode.sys(0, c(vec.tail.x[i],vec.tail.y[i]), params)[[1]]
		vec.dir.x[i] <- dir[1]
		vec.dir.y[i] <- dir[2]
	}

	# draw
	dev.new()
	plot(vec.tail.x, vec.tail.y, type='n', xlab=xlabel, ylab=ylabel)
	arrow.plot(
		vec.tail.x, vec.tail.y,
		vec.dir.x, vec.dir.y,
		arrow.ex=.05, length=.05, col='blue', lwd=1,
		true.angle = TRUE)

	####### draw nullclines ######
	# compute nullcline
	x <- seq(xrange[1], xrange[2], xstep/10)
	dx.nc.y <- rep(0, length(x))
	dy.nc.y <- rep(0, length(x))
	
	for (i in 1:length(x)) {
		y <- nullcline(x[i], params)
		dx.nc.y[i] <- y[1]
		dy.nc.y[i] <- y[2]
	}
	
	# draw
	lines(x, dx.nc.y, col='red')
	lines(x, dy.nc.y, col='red')

	####### draw trajectories #######
	n.starts = length(x.ini)
	points(x.ini, y.ini)
	for (i in 1:n.starts) {
		traj <- ode(y = c(y1=x.ini[i],y2=y.ini[i]),
				func = ode.sys,
				times = times, 
				parms = params,
				events=list(func=ode.event, root=TRUE), rootfun=ode.root)
		lines(traj[,2], traj[,3], col='black')
	}
}

# draw trajectory over time
# input:
#	ode.sys: system of ODEs
#	params: parameters for ode.sys
#	x.ini, y.ini: initial condition
#	times: time steps 
#	xlabel, ylabel: names of the first and the second depedent variables
draw.trajectory.over.time <- function(ode.sys, ode.root, ode.event, params, 
										x.ini, y.ini, times,
										xlabel, ylabel) {
	for (i in 1:length(x.ini)) {
		traj <- ode(y = c(y1=x.ini[i],y2=y.ini[i]),
				func = ode.sys,
				times = times, 
				parms = params,
				events=list(func=ode.event, root=TRUE), rootfun=ode.root)
				colnames(traj) <- c("t", xlabel, ylabel)
		dev.new()
		plot(traj, col='black')
	}
}