mpow <- function(M,k) {
  ## calculates matrix powers using a simple eigen-value decomposition
  if (k == 0) return (diag(1,ncol(M)))
  x <- eigen(M)
  U <- x$vectors
  L <- diag(x$values)
  return (U %*% L^k %*% solve(U))
}

## Simulating the evolution of a nucleotide
JC <- function(alpha) {
  ## construct Jukes-Cantor substitution matrix
  P <- matrix(0,4,4)
  for (i in 1:4) {
    for (j in 1:4)
      P[i,j] <- alpha
    P[i,i] <- 1-3*alpha
  }
  P
}

sample.nuc <- function() {
  sample(4,1)                         #sample uniformly from 1:4
}

evolve.nuc <- function(X,alpha,g) {
  ## get subtitution matrix
  P <- JC(alpha)
  ## get generation g probability vector
  pi.g <- mpow(P,g)[X,]
  ## sample Y
  sample(4,1,prob=pi.g)
}

## Simulating the evolution of a nucleotide sequence
sample.seq <- function(n) {
  sample(4,n,replace=TRUE)
}

evolve.seq <- function(X,alpha,g) {
  ## get transition probablity matrix
  P <- JC(alpha)
  P.g <- mpow(P,g)
  ## define single nuc. evolution 
  evolve <- function(X) {
    pi.g <- P.g[X,]
    sample(4,1,prob=pi.g)
  }
  ## evolve each nucleotide
  sapply(X,evolve)
}

## Calculating the likelihood
full.lhd <- function(alpha,X,Y,g) {
  P <- mpow(JC(alpha),g)
  nuc.prob <- function(nx,ny) P[nx,ny]
  factors <- mapply(nuc.prob,X,Y)
  prod(factors)
}

make.lhd <- function(X,Y,g) function(alpha) full.lhd(alpha,X,Y,g)

full.log.lhd <- function(alpha,X,Y,g) {
  P <- mpow(JC(alpha),g)
  nuc.prob <- function(nx,ny) log(P[nx,ny])
  terms <- mapply(nuc.prob,X,Y)
  sum(terms)
}

make.log.lhd <-
  function(X,Y,g) function(alpha) full.log.lhd(alpha,X,Y,g)


## difference based...
full.log.lhd <- function(alpha,X,Y,g) {
  no.equal <- sum(X == Y)
  no.diff <- sum(X != Y)
  P <- mpow(JC(alpha),g)
  prob.equal <- P[1,1]
  prob.diff <- P[1,2]
  no.equal * log(prob.equal) + no.diff * log(prob.diff)
}

## plotting
g <- 5
true.alpha <- 0.05
X <- sample.seq(1000)
Y <- evolve.seq(X,true.alpha,g)
log.lhd <- make.log.lhd(X,Y,g)

alphas <- seq(0,1/3,by=0.01)
plot(alphas,sapply(alphas,log.lhd),type='l',col='blue',
     main='Likelihood function',
     ylab=expression(loglhd(alpha)), xlab=expression(alpha))
abline(v=true.alpha, col='darkgreen', lty='dashed')
mtext('True value',3,at=true.alpha,col='darkgreen')


## maximum likelihood using nlm
ML.alpha <- nlm(function(alpha) -log.lhd(alpha), 0.1)