library(hdrm)
library(flsa)
library(neariso)

# Toy example
y <- c(0,0,0,1,1,1,0,0,0)
plot(y, pch=19, bty='n', las=1)

# Objective function; note that y is a better solution, but coordinate descent
# is stuck at 0
Q <- function(theta, y, lam=0.5) {
  n <- length(y)
  D <- bandSparse(n-1, n, c(0,1), list(rep(1, n-1), rep(-1, n-1)))
  drop(1/2*crossprod(y-theta) + lam*sum(abs(D%*%theta)))
}
Q(rep(0,9), y)
Q(y, y)

# Homegrown ADMM
admm <- function(y, lam, rho=0.1, iter=10) {
  n <- length(y)
  D <- bandSparse(n-1, n, c(0,1), list(rep(1, n-1), rep(-1, n-1)))
  tht <- rep(0, n)
  del <- rep(0, n-1)
  u <- rep(0, n-1)
  for (i in 1:iter) {
    tht <- solve(rho*crossprod(D) + diag(n)) %*% (y + rho*crossprod(D, del-u))
    del <- hdrm:::soft(rho*(D%*%tht + u), lam)/rho
    u <- u + D%*%tht - del
  }
  list(theta=tht, delta=del, u=u)
}
admm(y, 1/2, rho=0.5, iter=1)
admm(y, 1/2, rho=0.5, iter=2)
admm(y, 1/2, rho=0.5, iter=3)
res <- admm(y, 1/2, rho=0.5, iter=20)

# Note that this solutions is better than either 0 or y from before
Q(res$theta, y)

# ADMM for toy data
plot(y, pch=19, bty='n', las=1)
lines(1:10-0.5, rep(0, 10), type='s')
text(8, 0.9, 'Initial value')

for (i in 1:5) {
  plot(y, pch=19, bty='n', las=1)
  tht <- as.numeric(admm(y, 1/2, rho=0.5, iter=2*i)$theta)
  lines(1:10-0.5, c(tht, tht[9]), type='s')
  text(8, 0.9, paste('Iteration', 2*i))
}

# flsa for toy data
flsa(y, lambda1=0, lambda2=1/2)

# Fused lasso: CGH data
l2r <- readData(glioma)
fit <- flsa(l2r)
l2 <- c(0.1, 0.5, 1, 3)
plot(l2r, pch=16, cex=0.7, xlab='Genome order', ylab=expression(Log[2]*' ratio'), bty='n', las=1)
for (i in 1:length(l2)) {
  plot(l2r, pch=16, cex=0.7, xlab='Genome order', ylab=expression(Log[2]*' ratio'), col='gray', bty='n', las=1)
  lines(1:length(l2r), flsaGetSolution(fit, lambda1=0.1, lambda2=l2[i]), col='slateblue', lwd=3)
  abline(h=0, lty=2)
  mtext(bquote(lambda[2] == .(l2[i])))
}

# Fused lasso: Image de-noising
X <- matrix(0, 100, 100)
X[1:40,1:40] <- 1
X[61:100,1:40] <- 1
X[1:40,61:100] <- 1
X[61:100,61:100] <- 1
Y <- X + rnorm(100*100)
if (file.exists('cache/flsa-2d.rds')) {
  fit <- readRDS('cache/flsa-2d.rds')
} else {
  fit <- flsa(Y)
  saveRDS(fit, 'cache/flsa-2d.rds')
}
l2 <- c(0.1, 0.5, 1, 2, 5)
for (i in 1:length(l2)) {
  Z <- matrix(flsaGetSolution(fit, lambda2=l2[i]), 100, 100)
  par(mfrow=c(1,3), mar=c(0, 0, 1.5, 0))
  image(X, col=gray.colors(99), xaxt='n', yaxt='n'); mtext('True')
  image(Y, col=gray.colors(99), xaxt='n', yaxt='n'); mtext('Noisy')
  image(Z, col=gray.colors(99), xaxt='n', yaxt='n')
  mtext(bquote("FL: "*lambda[2] == .(l2[i])))
}

# Nearly isotonic regression
DF <- read.delim('https://s3.amazonaws.com/pbreheny-data-sets/warming.txt')
y <- DF$Raw + runif(nrow(DF), -0.005, 0.005)
res <- neariso(y)
lam <- c(0.5, 2)
plot(DF$Year, y, pch=19, xlab='Year', ylab='LOTI', bty='n', las=1)
for (i in 1:length(lam)) {
  plot(DF$Year, y, pch=19, col='gray', xlab='Year', ylab='LOTI', bty='n', las=1)
  lines(DF$Year, nearisoGetSolution(res, lam[i])$beta, lwd=3, col='red')
  mtext(bquote(lambda[2] == .(l2[i])))
}