## Simulation of the Dual Constraint model of the Neuromuscular Junction
## 
## This code reproduces the data in Figure 10.17 of "Principles of
## Computational Modelling in Neuroscience" by Sterratt, Graham,
## Gillies and Willshaw.
## 
## Please report any problems to David Sterratt
## <david.c.sterratt@ed.ac.uk>

## The deSolve package is required to run the simulations
## If it is not installed, install it.
if (!require("deSolve")) {
  install.packages("deSolve")
  require("deSolve")
}

##' This function simulates dual constraint model of NMJ. The full
##' details of the model can be found in Section 10.6.3 of Principles
##' of Computational Modelling in Neuroscience.
##'
##' @title Function to simulate dual constraint model of NMJ
##' @param N Number of axons
##' @param M Number of muscle fibres
##' @param A0 Total amount of presynaptic substance
##' @param alpha Forward rate constant of binding reaction
##' @param beta Backward rate constant of binding reaction 
##' @param gamma Forward rate constant of transport reaction
##' @param delta Backward rate constant of transport reaction
##' @param theta Threshold below which terminals are unviable
##' @param T.max Time of simulation
##' @param method Method used to solve ODE. See documentation of ode
##' for more details
##' @param ... Other parameters to pass to ode
##' @return A list containing
##' - Cnm Values over time of the factor C at the terminal of axon n
##'   at endplate m
##' - Anm Values over time of the factor A at the terminal of axon n
##'   at endplate m
##' - time A vector of time points at which Anm and Cnm were measured
##' - pars A list of the input parameters
##' @author David Sterratt
nmj <- function(N=6, M=240,
                A0=80, alpha=45, beta=0.4, gamma=3, delta=2, theta=0.01,
                T.max=0.1, method="lsoda", ...) {
  
  ## Construct a list of parameters from the inputs to pass to the
  ## model function later. We fix B0 to 1 without loss of generality.
  pars <- list(M  = M,
               N  = N,
               alpha = alpha,
               beta = beta,
               gamma = gamma,
               delta = delta,
               theta = theta,
               B0 = 1,
               A0 = A0,
               T.max = T.max)
  
  ## Model function of NMJ. This returns the derivatives dCnm/dt and
  ## dAnm/dt as a function of the parameters (Pars) and the Anm and
  ## Cnm (collectively given in the State argument).
  NMJmod <- function(Time, State, Pars) {
    with(Pars, {
      ## Read the State vector into Cnm and Anm. The first M*N
      ## elements of the State are Cnm and the next N*M elements are
      ## Anm. Cnm and Anm are both N by M matrices.
      Cnm <- matrix(State[1:(M*N)], N, M)
      Anm <- matrix(State[M*N+(1:(M*N))], N, M)

      ## Compute the amount of presynaptic factor left in each axon
      An <- A0 - rowSums(Cnm) - rowSums(Anm)

      ## Perform some sanity checks on Anm
      if (any(Anm<0)) {
        stop("An Anmm  is less than zero")
      }
      if (any(Cnm<0)) {
        stop("A Cnmm  is less than zero")
      }

      ## Compute the amount of postsynaptic factor left at each endplate
      Bm <- B0 - colSums(Cnm)
      if (any(Bm<0)) {
        stop("A Bm  is less than zero")
        Bm[Bm<0] <- 0
      }

      ## This vector indicates if the terminal is viable
      nunm <- Cnm > theta

      ## This computes the total number of viable terminals per axon
      nun <- rowSums(nunm)

      ## Compute the derivatives
      dCnm <- (alpha*Anm*matrix(Bm, N, M, byrow=TRUE) - beta)*Cnm
      dAnm <- nunm*gamma*matrix(An/(nun+1E-4), N, M) - delta*Anm/(Cnm+1E-7)
      return(list(c(dCnm, dAnm)))
    })
  }

  ## Now that the model function has been defined, specify the initial
  ## conditions. The Cnm are random numbers between 0 and 1/6, and the
  ## Anm variables are all zero.
  State0 <- c(runif(M*N, 0, 1/6), rep(0, M*N))

  ## Specify times at which to read out variables. Regardless of T.max
  ## there will be 100 time points.
  times <- seq(0, T.max, len=100)

  ## Solve the system of ODEs defined by the model function, given the
  ## initial conditions
  out <- as.data.frame(ode(func=NMJmod, y=State0,
                           parms=pars, times=times,
                           method=method, ...))

  ## Rearrange time column data into N-by-M-by-100 3D arrays. 
  Cnm <- array(as.matrix(t(out[,2:(1+N*M)])),         c(N, M, nrow(out)))
  Anm <- array(as.matrix(t(out[,(2+N*M):(1+2*N*M)])), c(N, M, nrow(out)))

  ## Return the parameters and data
  return(list(pars=pars, time=out[,"time"], Cnm=Cnm, Anm=Anm, N=N, M=M,
              A0=A0, B0=1 ))
}

## 

##' From http://cran.r-project.org/doc/contrib/R-and-octave-2.txt.
##'
##' <details>
##' @title 
##' @param a 
##' @param b 
##' @return 
##' @author David Sterratt
meshgrid <- function(a,b) {
  return(list(x=outer(b*0, a, FUN="+"),
              y=outer(b, a*0, FUN="+")))
}

##' Create a Willshaw plot of a weight matrix.
##'
##' @title Willshaw plot
##' @param w Weight matrix
##' @param add If TRUE, add to an existing plot. If FALSE, create a
##' new plot.
##' @param s Maximum fraction of unit square occupied by square
##' representing weight. 
##' @param ylab Label of y-axis.
##' @param xlab Label of x-axis.
##' @param bty Box type of axis. See documentation of par for more help.
##' @param xaxt Whether to plot x-axis tickmarks. See documentation of
##' par for more help.
##' @param yaxt Whether to plot y-axis tickmarks. See documentation of
##' par for more help.
##' @param ... 
##' @author David Sterratt
willshaw.plot <- function(w, add=FALSE, s=0.9, ylab="", xlab="", bty="o",
                        xaxt="n", yaxt="n",...) {
  ## Extract the number of rows and columns of w
  n <- nrow(w)
  m <- ncol(w)

  ## Create a list of the index pairs of every element of w
  mg <- meshgrid(1:m, 1:n)

  ## Create a new plot if required
  if(!add) {
    plot(NA, NA, xlim=c(1-1/2,m+1/2), ylim=rev(c(1-1/2,n+1/2)),
         ylab=ylab, xlab=xlab, bty=bty, xaxt=xaxt, yaxt=yaxt, ...)
  }

  ## Plot the rectangles
  x <- matrix(mg$x, ncol=1) 
  y <- matrix(mg$y, ncol=1) 
  rect(x-s*0.5*sqrt(as.vector(w)), y-s*0.5*sqrt(as.vector(w)),
       x+s*0.5*sqrt(as.vector(w)), y+s*0.5*sqrt(as.vector(w)), col="black")
}

##' This function plots the results of an NMJ simulation produced by
##' the nmj() function. 
##'
##' @title Plot the results of an NMJ simulation
##' @param out A list of variables produced by the nmj() function.
##' @param cols An array of colours in which to plot the values of mu_m
##' @author David Sterratt
plot.sim <- function(out, cols="black") {

  ## Set graphics parameters. See help for par for more details.
  par(mfcol=c(4,2))                     # 4 by 2 grid arrangement
  par(cex=1)                            # Expansion of characters
  par(oma=c(2, 0, 0, 0))                # Outer margin
  par(bty="l")                          # The box type of the axes
  par(xpd=NA)                           # Allow drawing outside the margine

  ## The with(out, ) idiom makes the members of the list out available
  with(out, {
    ## Compute An
    An <- pars$A0 - apply(Anm, c(1,3), sum) - apply(Cnm, c(1,3), sum)

    ## Compute indicator of presence of synapse
    nunm <- Cnm > pars$theta

    ## Compute motor unit size
    nun <- apply(nunm, c(1,3), sum)

    ## Compute endplate unit size
    mum <- apply(nunm, c(2,3), sum)

    ## Compute histogram of singly, double &c innervated endplates
    mum.hist <- matrix(0, 1+pars$N, ncol(mum))
    rownames(mum.hist) <- 0:pars$N
    for (i in 1:ncol(mum)) {
      tmum <- table(mum[,i])
      mum.hist[names(tmum),i] <- tmum
    }

    ## Plot left column - plots of Cnm at different times

    par(mar=c(1, 1, 1, 4))              # Set margin

    willshaw.plot(out$Cnm[,1:20,1])
    mtext(expression(paste(italic(t), "= 1")), line=0.3)
    willshaw.plot(out$Cnm[,1:20,10])
    mtext(expression(paste(italic(t), "= 10")), line=0.3)
    willshaw.plot(out$Cnm[,1:20,50])
    mtext(expression(paste(italic(t), "= 50")), line=0.3)
    willshaw.plot(out$Cnm[,1:20,100])
    mtext(expression(paste(italic(t), "= 100")), line=0.3)
    
    ## Plot right column - statistics

    par(mar=c(1, 1, 1, 1))              # Set margin

    ## Time course of C_n1
    matplot(time, t(Cnm[,1,]), type="l", lty=1, col="black",
            ylab="", xaxt="n", xlab="",
            ylim=c(0,1), yaxp=c(0,1,2), bty="n")
    axis(1, labels=FALSE)
    mtext(expression(italic(C)[italic(nm)]), 2, 2.3, las=0)

    ## Time course of A_n in each axon
    matplot(time, t(An), type="l", lty=1, col="black",
            ylab="", xaxt="n", xlab="",
            ylim=c(0, A0), yaxp=c(0, A0, 2), bty="n")
    axis(1, labels=FALSE)
    mtext(expression(italic(A)[italic(n)]), 2, 2.3, las=0)

    ## Time course of vu_n in each axon
    matplot(time, t(nun), type="l", lty=1, col="black",
            ylab="", xaxt="n", xlab="",
            ylim=c(0, M), yaxp=c(0, M, 2), bty="n")
    axis(1, labels=FALSE)
    mtext(expression(italic(nu)[italic(n)]), 2, 2.3, las=0)

    ## Time course of mu_m the amount of single, double &c innervation
    matplot(time, t(mum.hist), type="l", lty=1, col=cols,
            ylab="", xlab=expression(italic(t)),
            ylim=c(0,M), yaxp=c(0, M, 2), bty="n")
    mtext(expression(italic(mu)[italic(m)]), 2, 2.3, las=0)
  })
}

## Run the simulation
out <- nmj(T.max=30, hini=0.002, theta=0.01, method="euler")

## Plot the simulation
message("Running simulation...")
plot.sim(out, cols=rainbow(out$pars$N+1))