.packageName <- "ouch"
# This file is part of the OUCH package.
# Author: Aaron A. King <king at tiem dot utk dot edu>
# It is distributed under the GNU Public License (see the file GPL
# included)
# The OUCH package is maintained at
#          http://www.tiem.utk.edu/~king/ouch.
#
badness <- function (alpha, data, parsed.tree) {
  a <- exp(alpha);
  n <- length(data);
  w <- weight.matrix(a, parsed.tree);
  v <- scaled.covariance.matrix(a, parsed.tree);
  g <- glssoln(w,data,v);
  e <- g$residuals;
  sigmasq <- (e %*% solve(v,e)) / n;
  dim(sigmasq) <- 1;
  u <- n * (1 + log(2*pi*sigmasq)) + log(det(v));
  dim(u) <- 1;
  return(u);
}

branch.times <- function (topology, times) {

  term <- terminal.twigs(topology, times);
  N <- length(term);
  T <- max(times);

  bt <- matrix(data=0,nrow=N,ncol=N);

  bt[1,1] <- T;
  for (i in 2:N) {
    pedi <- pedigree(topology,term[i]);
    for (j in 1:(i-1)) {
      pedj <- pedigree(topology,term[j]);
      for (k in 1:length(pedi)) {
        if (any(pedj == pedi[k])) break;
      }
      bt[j,i] <- bt[i,j] <- times[pedi[k]];
    }
    bt[i,i] <- T;
  }
  return(bt);
}

brown.dev <- function(n = 1, topology, times, sigma, theta) {
  pt <- parse.tree(topology,times);
  v <- pt$branch.times;
  x <- rmvnorm(n, rep(theta,dim(v)[1]), as.numeric(sigma^2)*v);
  return(data.frame(x))
}


brown.fit <- function (data, topology, times) {
  pt <- parse.tree(topology,times);
  n <- pt$N;
  v <- pt$branch.times;
  w <- matrix(data=1,nrow=pt$N,ncol=1);
  dat <- data[!is.na(data)];
  g <- glssoln(w,dat,v);
  theta <- g$coeff;
  e <- g$residuals;
  sigma <- sqrt((e %*% solve(v,e))/n);
  dim(sigma) <- 1;
  u = n * (1 + log(2*pi*sigma*sigma)) + log(det(v));
  dim(u) <- 1;
  df <- 2;
  return(list(sigma=sigma,theta=theta,u=u,aic=u+2*df,sic=u+log(n)*df,df=df));
}

epochs <- function (topology, times, term) {
  N <- length(term);
  e <- vector(length=N,mode="list");
  for (k in 1:N) {
    p <- pedigree(topology,term[k]);
    e[[k]] <- times[p];	
  }
  return(e);
}

glssoln <- function(a, x, v, tol = 1e-12) {
  n <- length(x);
  vh <- t(chol(v));
  s <- svd(forwardsolve(vh,a));
                                        #   Can we be certain that the singular values are sorted in
                                        #   decreasing order?  (Probably not)
                                        #   k <- order(s$d,decreasing=T);
  svals <- s$d[s$d > tol * max(s$d)];
  r <- length(svals);
  svals <-  diag(1/svals,nrow=r,ncol=r);
  y <- (s$v[,1:r] %*% (svals %*% t(s$u[,1:r]))) %*% forwardsolve(vh,x);
  e <- a %*% y - x;
  dim(y) <- dim(y)[1];
  dim(e) <- n;
  return(list(coeff=y,residuals=e));
}

hansen.dev <- function(n = 1, topology, times, regimes, alpha, sigma, theta) {
  pt <- parse.tree(topology,times,regimes);
  w <- weight.matrix(alpha, pt);
  v <- scaled.covariance.matrix(alpha, pt);
  x <- rmvnorm(n, as.vector(w %*% theta), as.numeric(sigma^2)*v);
  return(data.frame(x))
}
hansen.fit <- function (data, topology, times, regimes, guess=0, interval=c(0.001,20), tol=1e-12) {
  pt <- parse.tree(topology,times,regimes);
  n <- pt$N;
  dat <- data[!is.na(data)];
  r <- optimize(badness,interval=log(interval),
                lower=log(interval[1]),upper=log(interval[2]),
                tol=tol,maximum=F,dat,pt);
  alpha = exp(r$minimum);

  w <- weight.matrix(alpha, pt);
  v <- scaled.covariance.matrix(alpha, pt);
  g <- glssoln(w,dat,v);
  theta <- g$coeff;
  e <- g$residuals;
  sigma <- sqrt((e %*% solve(v,e))/n);
  dim(sigma) <- 1;
  u = r$objective;
  dim(u) <- 1;
  df <- pt$R+3;
  return(list(alpha=alpha,sigma=sigma,theta=theta,u=u,aic=u+2*df,sic=u+log(n)*df,df=df));
}

parse.tree <- function (topology, times, regime.specs=NULL) {
  if ((length(topology) != length(times)) 
       || ((!is.null(regime.specs))
          && (length(topology) != length(regime.specs)))
      ) {
    warning('invalid tree')
    return(NULL);
  }
  term <- terminal.twigs(topology,times);
  N <- length(term);
  T <- max(times);
  bt <- branch.times(topology,times);
  e <- epochs(topology,times,term);
  if (is.null(regime.specs)) {
    pt <- list(N=N,T=T,term=term,branch.times=bt,epochs=e);
  } else {
    reg <- set.of.regimes(topology,regime.specs);    
    R <- length(reg);
    b <- regimes(topology, times, regime.specs, term);
    pt <- list(N=N,R=R,T=T,term=term,branch.times=bt,epochs=e,regime.set=reg,beta=b);
  }
  return(pt);
}

pedigree <- function (topology, k) {
  p <- k;
  k <- topology[k];
  while (k != 0) {
    p <- c(p, k);
    k <- topology[k];
  }
  return(p);
}

regimes <- function (topology, times, regime.specs, term) {
  N <- length(term);
  reg <- set.of.regimes(topology,regime.specs);
  R <- length(reg);
  beta <- vector(R*N, mode="list");
  for (i in 1:N) {
    for (k in 1:R) {
      p <- pedigree(topology, term[i]);
      n <- length(p);
      beta[[i + N*(k-1)]] <- as.integer(regime.specs[p[1:(n-1)]] == reg[k]);
    }
  }    
  return(beta);
}

rmvnorm <- function (n = 1, mu, sigma, tol = 1e-06) {
  p <- length(mu);
  if (!all(dim(sigma) == c(p,p)))
    stop("incompatible arguments");
  cf <- chol(sigma,pivot=F);
  X <- matrix(mu,n,p,byrow=T) + matrix(rnorm(p*n),n) %*% cf;
  if (n == 1) {
    return(drop(X));
  }  else {
    return(X);
  }
}

scaled.covariance.matrix <- function (alpha, parsed.tree) {
  T <- parsed.tree$T;
  bt <- parsed.tree$branch.times;			 
  if (alpha == 0) {
    V <- bt;
  } else {
    a <- 2*alpha;
    V <- exp(-a*T) * expm1(a*bt) / a;
  }
}

set.of.regimes <- function (topology, regime.specs) {
  n <- length(regime.specs);
  id <- seq(1,n)[topology > 0];       # find all non-root nodes
  reg <- sort(unique(regime.specs[id]));
  return(reg);		  
}

terminal.twigs <- function (topology, times) {
  n <- length(topology);
  return(seq(1,n)[times == max(times)]);
}

tree.plot <- function (topology, times, names = NULL, regimes = NULL) {
  rx <- range(times,na.rm=T);
  rxd <- 0.1*diff(rx);

  if (is.null(regimes))
    regimes <- factor(rep(1,length(topology)));

  levs <- levels(as.factor(regimes));
  palette <- rainbow(length(levs));

  for (r in 1:length(levs)) {
    y <- tree.layout(topology);
    x <- times;
    f <- which(topology > 0 & regimes == levs[r]);
    pp <- topology[f];
    X <- array(data=c(x[f], x[pp], rep(NA,length(f))),dim=c(length(f),3));
    Y <- array(data=c(y[f], y[pp], rep(NA,length(f))),dim=c(length(f),3));
    oz <- array(data=1,dim=c(2,1));
    X <- kronecker(t(X),oz);
    Y <- kronecker(t(Y),oz);
    X <- X[2:length(X)];
    Y <- Y[1:(length(Y)-1)];
    C <- rep(palette[r],length(X));
    if (r > 1) par(new=T);
    par(yaxt='n')
    plot(X,Y,type='l',col=C,xlab='time',ylab='',xlim = rx + c(-rxd,rxd),ylim=c(0,1));
    if (!is.null(names))
      text(X[seq(1,length(X),6)],Y[seq(1,length(Y),6)],names[f],pos=4);
  }
}

tree.layout <- function (topology) {
  root <- which(topology==0);
  return(arrange.tree(root,topology));
}

arrange.tree <- function (root, topology) {
  k <- which(topology==root);
  n <- length(k);
  reltree <- rep(0,length(topology));
  reltree[root] <- 0.5;
  p <- list(NULL);
  if (n > 0) {
    m <- rep(0,n);
    for (j in 1:n) {
      p[[j]] <- arrange.tree(k[j],topology);
      m[j] <- length(which(p[[j]] != 0));
    }
    cm <- c(0,cumsum(m));
    for (j in 1:n) {
      reltree <- reltree + (cm[j]/sum(m))*(p[[j]] != 0) + (m[j]/sum(m))*p[[j]];
    }
  }
  return(reltree);
}
weight.matrix <- function (alpha, parsed.tree) {
  N <- parsed.tree$N;
  R <- parsed.tree$R;
  T <- parsed.tree$T;
  ep <- parsed.tree$epochs;
  beta <- parsed.tree$beta;
  W <- matrix(data=0,nrow=N,ncol=R+1);
  W[,1] <- exp(-alpha*T);      
  for (i in 1:N) {
    delta <- diff(exp(alpha*(ep[[i]]-T)));
    for (k in 1:R) {
      W[i,k+1] <- -sum(delta * beta[[i+N*(k-1)]]);
    }
  }
  return(W);
}

# This file is part of the OUCH package.
# Author: Aaron A. King <king at tiem dot utk dot edu>
# It is distributed under the GNU Public License (see the file GPL
# included)
# The OUCH package is maintained at
#          http://www.tiem.utk.edu/~king/ouch.
#
