##############################################################
##              Plots                                       ##
##############################################################

# Scree plots (on a log scale)
my_ruv_scree <-
function(Y, Z=1, xlab="K", ...)
{
  residop = function(A,B) return(A - B%*%solve(t(B)%*%B)%*%t(B)%*%A) # residual operator
  m = nrow(Y)
  if (is.null(Z))
  {
    q = 0
  }
  else if (length(Z) == 1) {if (Z==1)  
  {
    Z = matrix(1,m,1)
    q = 1
  }}
  else 
  {
    q = ncol(Z)
  }
  if (q > 0)
  {
    Y = residop(Y,Z)
  }
  plotval = log(svd(Y)$d[1:(m-q)])
  plot(plotval,xlab=xlab,ylab="Log Eigenvalue",...)
}

# RLE plots
my_ruv_rle <-
function(Y, XZW=NULL, coloring=1, xlab="", ylims=c(-.5,.5), ...)
{
  residop = function(A,B) return(A - B%*%solve(t(B)%*%B)%*%t(B)%*%A) # residual operator
  if (!is.null(XZW)) Y = t(residop(Y, XZW))
  else Y = t(Y)
  medians = apply(Y,1,median)
  Y = Y-medians
  boxplot(Y,col=coloring,xaxt="n",outline=F,ylim=ylims,range=0,...)
  lines(c(0,ncol(Y)),c(0,0))
}

# Canonical Correlation Plots
my_ruv_cancorplot <-
function(X, W1, W2=NULL, col=c(1,3), xlab="K", ylab="cancor(X,W)^2", ...)
{
  K = ncol(W1)
  cc1 = cc2 = rep(0,K)
  for (k in 1:K)
  {
    cc1[k] = cancor(X,W1[,1:k,drop=FALSE])$cor[1]^2
    if (!is.null(W2)) cc2[k] = cancor(X,W2[,1:k,drop=FALSE])$cor[1]^2
  }
  plot(1:K,cc1,col=col[1],xlab=xlab,ylab=ylab, ylim=c(0,1), ...)
  if (!is.null(W2)) points(1:K,cc2,col=col[2], ...)
}

# P-value Histograms
my_ruv_pvalhist <-
function(pvals, l=1, xlab="", ...)
{
  if (is.matrix(pvals)) pvals = pvals[l,]
  hist(pvals,xlab=xlab,breaks=c(0,0.001,0.01,0.05,.1,.2,.3,.4,.5,.6,.7,.8,.9,1),main="",...)
}

## The next several plot routines all make a call to "my_ruv_niceplot."
## my_ruv_niceplot is basically just a nice version of "plot" that makes it easier to specify the colors and symbols (including text).

# P-value Plots
my_ruv_pvalplot <-
function(pvals, l=1, logged=FALSE, coloring=NULL, coloring_order=NULL, symbols=NULL, texts=NULL, small_cutoff=-Inf, ...)
{
  if (is.matrix(pvals)) pvals = pvals[l,]
  n = length(pvals)
  x=1:n
  neworder = order(pvals)
  y = pvals[neworder]
  if (!is.null(coloring)) coloring = coloring[neworder]
  if (!is.null(symbols)) symbols = symbols[neworder]
  if (!is.null(texts)) texts = texts[neworder]
  if (logged)
  {
    x = log(x)
    y = log(y)
    # use the small_cutoff when some of the p-values are absurdly small (or 0), and it messes up the plot.
    x = x[y > small_cutoff]
    coloring = coloring[y > small_cutoff]
    if (!is.null(symbols)) symbols = symbols[y > small_cutoff]
    if (!is.null(texts)) texts = texts[y > small_cutoff]
    y = y[y > small_cutoff]
    xlabel = "Log Rank"
    ylabel = "Log P-value"
    my_ruv_niceplot(x,y,coloring=coloring,coloring_order=coloring_order,symbols=symbols,texts=texts,xlab=xlabel,ylab=ylabel,...)
  }
  else 
  {
    xlabel = "Rank"
    ylabel = "P-value"
    my_ruv_niceplot(x,y,coloring=coloring,coloring_order=coloring_order,symbols=symbols,texts=texts,xlab=xlabel,ylab=ylabel,ylim=c(0,1),...)
  }
}

# Projection Plots
my_ruv_projectionplot <-
function(byx, alpha, bwx, l=1, coloring=NULL, coloring_order=NULL, symbols=NULL, texts=NULL, ...)
{
  if (is.list(alpha)) 
  {
    if (is.null(alpha$projectionplotalpha)) alpha = alpha$alpha
    else alpha = alpha$projectionplotalpha
  }
  y = as.vector(byx[l,])
  graddir = bwx[l,] / sqrt(sum(bwx[l,]^2))
  gradslope = sum(bwx[l,]*graddir)
  x = as.vector(t(matrix(graddir))%*%alpha)
  linex = c(min(x),max(x))
  liney = linex*gradslope
  xlabel = "Gradient Factor"
  my_ruv_niceplot(x,y,coloring=coloring,coloring_order=coloring_order,symbols=symbols,texts=texts,xlab=xlabel,ylab=expression(b[YX]), ...)
  lines(linex,liney,col="black")
}

my_ruv_factor_projectionplot <-
function(fullfit, k, factor, l=1, coloring=NULL, coloring_order=NULL, symbols=NULL, texts=NULL, ...)
{
  if (k==0) keepindex = NULL
  else keepindex = 1:k
  if (factor <= k) keepindex = keepindex[-factor]
  x = fullfit$alpha[factor,]
  y = as.vector(fullfit$byx[l,])
  if (length(keepindex)>0) 
  {
    if (is.null(fullfit$projectionplotalpha)) y = y - as.vector(fullfit$bwx[l,keepindex,drop=FALSE]%*%fullfit$alpha[keepindex,,drop=FALSE])
    else y = y - as.vector(fullfit$bwx[l,keepindex,drop=FALSE]%*%fullfit$projectionplotalpha[keepindex,,drop=FALSE])
  }
  slope = fullfit$bwx[l,factor]
  linex = c(min(x),max(x))
  liney = linex*slope
  xlabel = paste("Factor", factor)
  ylabel = bquote(Adjusted~b[YX]~~(K~"="~.(k)))
  my_ruv_niceplot(x,y,coloring=coloring,coloring_order=coloring_order,symbols=symbols,texts=texts,xlab=xlabel,ylab=ylabel, ...)
  lines(linex,liney,col="black")
}

my_ruv_topprojectionplot <-
function(byx, alpha, bwx, pvals, l=1, cutoffs = c(15, 40, 75, 150, 300), topcolors = c("blueviolet", "royalblue", "seagreen1", "orange", "orange4"), rankset=NULL, ...)
{
  if (is.matrix(pvals)) pvals = pvals[l,]
  n = length(pvals)
  if (is.null(rankset)) rankset = rep(TRUE,n)
  neworder = order(pvals)
  neworder = neworder[rankset[neworder]]
  cutoffs = cutoffs[cutoffs < length(neworder)]
  topcolors = topcolors[1:length(cutoffs)]
  coloring = rep(NA, n)
  texts = rep(NA, n)
  for (i in length(topcolors):1) coloring[neworder[1:cutoffs[i]]] = topcolors[i]
  texts[neworder[1:cutoffs[1]]] = as.character(1:cutoffs[1])
  my_ruv_projectionplot(byx, alpha, bwx, l=l, coloring=coloring, texts=texts, symbols=20,...)
}

# Variance Plots
my_ruv_varianceplot <- 
function(fit, varbetahat=NULL, trans=my_fourthroot, l=1, coloring=NULL, coloring_order=NULL, symbols=NULL, texts=NULL, xlab="trans(sigma^2 hat)", ylab="trans(betahat^2)", ...)
{
  x = fit$sigma2
  y = as.vector(fit$betahat[l,]^2)
  neworder = order(x)
  if (!is.null(coloring)) coloring = coloring[neworder]
  if (!is.null(symbols)) symbols = symbols[neworder]
  if (!is.null(texts)) texts = texts[neworder]
  x = trans(x[neworder]) # other transformations are possible.  see below.
  y = trans(y[neworder])
  my_ruv_niceplot(x, y, coloring=coloring, coloring_order=coloring_order, symbols=symbols, texts=texts, xlab=xlab, ylab=ylab, ...)
  if (!is.null(varbetahat))
  {
    varbetahat = trans(varbetahat[neworder])
    lines(x,varbetahat,col="black")
  }
}

my_fourthroot = function(x) # for use with variance plots
{
  return(x^(1/4))
}

my_identitytrans = function(x) # for use with variance plots
{
  return(x)
}

my_ruv_topvarianceplot <-
function(fit, pvals, varbetahat=NULL, trans=my_fourthroot, l=1, cutoffs = c(15, 40, 75, 150, 300), topcolors = c("blueviolet", "royalblue", "seagreen1", "orange", "orange4"), xlab="trans(sigma^2 hat)", ylab="trans(betahat^2)", rankset=NULL, ...)
{
  if (is.matrix(pvals)) pvals = pvals[l,]
  n = length(pvals)
  if (is.null(rankset)) rankset = rep(TRUE,n)
  neworder = order(pvals)
  neworder = neworder[rankset[neworder]]
  cutoffs = cutoffs[cutoffs < length(neworder)]
  topcolors = topcolors[1:length(cutoffs)]
  coloring = rep(NA, n)
  texts = rep(NA, n)
  for (i in length(topcolors):1) coloring[neworder[1:cutoffs[i]]] = topcolors[i]
  texts[neworder[1:cutoffs[1]]] = as.character(1:cutoffs[1])
  my_ruv_varianceplot(fit, varbetahat=varbetahat, l=l, trans=trans, coloring=coloring, texts=texts, symbols=20, xlab=xlab, ylab=ylab, ...)
}

# PC plots
my_ruv_pcgridplot <-
function(W1, W2=NULL, kset=c(1,2,3), coloring=NULL, coloring_order=NULL, symbols=NULL, texts=NULL, symbol_overwrite=TRUE, ...)
{
  kn = length(kset)
  par(mfrow=c(kn,kn))
  for (i in 1:kn)
  {
    for (j in 1:kn)
    {
      if (i == j) 
      {
        if (i == 1) my_controlvsfulllabelplot(1)
        else if (i < kn) my_controlvsfulllabelplot(2)
        else my_controlvsfulllabelplot(3)
      }
      if (j > i)
      {
        my_ruv_niceplot(W1[,i], W1[,j], coloring=coloring, coloring_order=coloring_order, symbols=symbols, texts=texts, symbol_overwrite=symbol_overwrite, xlab=paste("Factor",i), ylab=paste("Factor",j), cex.lab=1.5, ...)
      }
      if (j < i)
      {
        if (!is.null(W2)) my_ruv_niceplot(W2[,j], W2[,i], coloring=coloring, coloring_order=coloring_order, symbols=symbols, texts=texts, symbol_overwrite=symbol_overwrite, xlab=paste("Factor",j), ylab=paste("Factor",i), cex.lab=1.5, ...)
        else frame()
      }
    }
  }
  par(mfrow=c(1,1))
}

my_controlvsfulllabelplot = function(mod) # for use with PC plots
{
  frame()
  if (mod==1 | mod==2)
  {
    arrows(.2, .75, .8, .75)
    text(x=.5, .8, labels="All genes", cex=2)
  }
  if (mod==2 | mod==3)
  {
    arrows(.8, .25, .2, .25)
    text(x=.5, .3, labels="Control genes", cex=2)
  }
}

# This is the function my_ruv_niceplot.  It is called by several of the routines above. 

my_ruv_niceplot <-
function(x,y,coloring=NULL,coloring_order=NULL,symbols=NULL,texts=NULL,symbol_overwrite=TRUE,...)
{
  n = length(y)
  plot(c(min(x), max(x)), c(min(y),max(y)), col=0, ...)
  if (is.null(coloring)) coloring = rep(1,n)
  if (length(coloring)==1) coloring = rep(coloring, n)
  coloring[is.na(coloring)] = 8
  if (is.null(coloring_order)) coloring_order = my_get_coloring_order(coloring)
  if (is.null(symbols)) symbols = rep(1,n)
  if (length(symbols)==1) symbols = rep(symbols, n)
  symbols[is.na(symbols)] = 1
  if (!is.null(texts))
  {
    textsx = x[!is.na(texts)]
    textsy = y[!is.na(texts)]
    textscoloring = coloring[!is.na(texts)]
    if (symbol_overwrite)
    {
      x = x[is.na(texts)]
      y = y[is.na(texts)]
      coloring = coloring[is.na(texts)]
      symbols = symbols[is.na(texts)]
    }
    texts = texts[!is.na(texts)]
  }
  cn = length(coloring_order)
  for (i in 1:cn)
  {
    if (sum(coloring==coloring_order[i])>0) points(x[coloring==coloring_order[i]], y[coloring==coloring_order[i]], col=coloring_order[i], pch=symbols[coloring==coloring_order[i]])
    if (!is.null(texts)) if (sum(textscoloring==coloring_order[i])>0) text(textsx[textscoloring==coloring_order[i]], textsy[textscoloring==coloring_order[i]], labels=texts[textscoloring==coloring_order[i]], col=coloring_order[i])
  }
}

my_get_coloring_order <- # The default coloring order, used by my_ruv_niceplot
function(coloring)
{
  if (is.numeric(coloring)) collevels = as.numeric(levels(as.factor(coloring)))
  else collevels = levels(as.factor(coloring))
  coln = length(collevels)
  colcounts = rep(0,coln)
  for (i in 1:coln) colcounts[i] = sum(coloring==collevels[i], na.rm=TRUE)
  neworder = order(colcounts)
  neworder = neworder[coln:1]
  return(collevels[neworder])
}

# Top Rank Count Plot
my_ruv_toprankcountplot <-
function(counttable, kset, thresholds)
{
  plot(c(min(kset),max(kset)),c(min(counttable),max(counttable)), xlab="K", ylab=c("Number of Top-Ranked positive controls", paste(thresholds, collapse=", ")), col=0)
  for (i in 1:ncol(counttable)) lines(kset,counttable[,i])
}




##############################################################
##              Tables and Counts                           ##
##############################################################

my_ruv2toprankcounttables <-
function(Y,X,ctl,pctl,Z=1,v=NULL,fullW=NULL,l=1,bin=10,thresholds=c(20,40,60,80,100),kset=c(0:20,2*(11:25),5*(11:20),10*(11:100)),rankset=NULL,inputcheck=TRUE)
{
  if (inputcheck) inputcheck1(Y,X,Z,ctl)
  Y = adjust_v(Y,v,ctl)
  m = nrow(Y)
  if (is.null(Z)) q = 0
  else if (length(Z) == 1) q = 1
  else q = ncol(Z)
  kset = kset[kset < m-q-1]
  kn = length(kset)
  threshn = length(thresholds)
  if (is.null(rankset)) rankset = rep(TRUE,ncol(Y))
  countmat = countmat.ebayes = countmat.rsvar = countmat.rsvar.ebayes = countmat.evar = countmat.beta = matrix(0,kn,threshn)
  fullW = RUV2(Y,X,ctl,0,Z=Z,v=v,inputcheck=FALSE)$fullW
  for (ki in 1:kn)
  {
    ruv2fit = RUV2(Y,X,ctl,kset[ki],Z=Z,v=v,fullW=fullW,inputcheck=FALSE)
    ruv2fit = variance_adjust(ruv2fit, bin=bin)
    countmat[ki, ] = as.vector(my_ruv_toprankcount(ruv2fit$p[l,rankset], pctl[rankset], thresholds))
    countmat.ebayes[ki, ] = as.vector(my_ruv_toprankcount(ruv2fit$p.ebayes[l,rankset], pctl[rankset], thresholds))
    countmat.rsvar[ki, ] = as.vector(my_ruv_toprankcount(ruv2fit$p.rsvar[l,rankset], pctl[rankset], thresholds))
    countmat.rsvar.ebayes[ki, ] = as.vector(my_ruv_toprankcount(ruv2fit$p.rsvar.ebayes[l,rankset], pctl[rankset], thresholds))
    countmat.evar[ki, ] = as.vector(my_ruv_toprankcount(ruv2fit$p.evar[l,rankset], pctl[rankset], thresholds))
    countmat.beta[ki, ] = as.vector(my_ruv_toprankcount(-abs(ruv2fit$betahat[l,rankset]), pctl[rankset], thresholds))
  }
  return(list(countmat=countmat, countmat.ebayes=countmat.ebayes, countmat.rsvar=countmat.rsvar, countmat.rsvar.ebayes=countmat.rsvar.ebayes, countmat.evar=countmat.evar, countmat.beta=countmat.beta, kset=kset, thresholds=thresholds))
}

my_ruv4toprankcounttables <-
function(Y,X,ctl,pctl,Z=1,v=NULL,fullW0=NULL,l=1,bin=10,thresholds=c(20,40,60,80,100),kset=c(0:20,2*(11:25),5*(11:20),10*(11:100)),rankset=NULL,inputcheck=TRUE)
{
  if (inputcheck) inputcheck1(Y,X,Z,ctl)
  Y = adjust_v(Y,v,ctl)
  m = nrow(Y)
  if (is.null(Z)) q = 0
  else if (length(Z) == 1) q = 1
  else q = ncol(Z)
  kset = kset[kset < m-q-1]
  kn = length(kset)
  threshn = length(thresholds)
  if (is.null(rankset)) rankset = rep(TRUE,ncol(Y))
  countmat = countmat.ebayes = countmat.rsvar = countmat.rsvar.ebayes = countmat.evar = countmat.beta = matrix(0,kn,threshn)
  fullW0 = RUV4(Y,X,ctl,0,Z=Z,v=v,inputcheck=FALSE)$fullW0
  for (ki in 1:kn)
  {
    ruv4fit = RUV4(Y,X,ctl,kset[ki],Z=Z,v=v,fullW0=fullW0,inputcheck=FALSE)
    ruv4fit = variance_adjust(ruv4fit, bin=bin)
    countmat[ki, ] = as.vector(my_ruv_toprankcount(ruv4fit$p[l,rankset], pctl[rankset], thresholds))
    countmat.ebayes[ki, ] = as.vector(my_ruv_toprankcount(ruv4fit$p.ebayes[l,rankset], pctl[rankset], thresholds))
    countmat.rsvar[ki, ] = as.vector(my_ruv_toprankcount(ruv4fit$p.rsvar[l,rankset], pctl[rankset], thresholds))
    countmat.rsvar.ebayes[ki, ] = as.vector(my_ruv_toprankcount(ruv4fit$p.rsvar.ebayes[l,rankset], pctl[rankset], thresholds))
    countmat.evar[ki, ] = as.vector(my_ruv_toprankcount(ruv4fit$p.evar[l,rankset], pctl[rankset], thresholds))
    countmat.beta[ki, ] = as.vector(my_ruv_toprankcount(-abs(ruv4fit$betahat[l,rankset]), pctl[rankset], thresholds))
  }
  return(list(countmat=countmat, countmat.ebayes=countmat.ebayes, countmat.rsvar=countmat.rsvar, countmat.rsvar.ebayes=countmat.rsvar.ebayes, countmat.evar=countmat.evar, countmat.beta=countmat.beta, kset=kset, thresholds=thresholds))
}

my_ruv_toprankcount <-
function(pvals, pctl, thresholds, l=1)
{
  if (is.matrix(pvals)) pvals = pvals[l,]
  threshn = length(thresholds)
  counts = matrix(0,1,threshn)
  colnames(counts) = thresholds
  rownames(counts) = "count"
  neworder = order(pvals)
  for (i in 1:threshn)
  {
    counts[i] = sum(pctl[neworder[1:thresholds[i]]])
  }
  countmat = matrix(counts,1,)
  colnames(countmat) = thresholds
  return(countmat)
}

my_ruv_topranktable <-
function(pvals, beta=NULL, rankby=NULL, info=NULL, N=40, l=1, rankset=NULL)
{
  if (is.matrix(pvals)) pvals = pvals[l,]
  n = length(pvals)
  if (is.null(rankset)) rankset = rep(TRUE,n)
  if (is.null(rankby)) neworder = order(pvals)
  else neworder = order(rankby)
  topindices = neworder[rankset[neworder]][1:N]
  adjp = p.adjust(pvals, method="BH")
  if (!is.null(beta)) beta = as.character(signif(beta,2))
  tableinfo = cbind(rank=1:N, cbind(p=as.character(signif(pvals,1)), BH=as.character(signif(adjp,1)), beta, info)[topindices,])
  return(tableinfo)
}