suppressPackageStartupMessages({
  library(hdrm)
  library(flsa)
  library(genlasso)
  library(CVXR)
})

# 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(y, lambda1 = 0, lambda2 = 1/2)

fit <- flsa(y)

flsaGetSolution(fit, lambda1 = 0.1, lambda2 = 1/2)

# Fused lasso: CGH data
fit <- flsa(glioma)
l2 <- c(0.1, 0.5, 1, 3)
plot(
  glioma,
  pch = 16,
  cex = 0.7,
  xlab = 'Genome order',
  ylab = expression(Log[2] * ' ratio'),
  bty = 'n',
  las = 1
)
for (i in 1:length(l2)) {
  plot(
    glioma,
    pch = 16,
    cex = 0.7,
    xlab = 'Genome order',
    ylab = expression(Log[2] * ' ratio'),
    col = 'gray',
    bty = 'n',
    las = 1
  )
  lines(
    1:length(glioma),
    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])))
}

# Global warming data (adding a tiny amount of noise)
df <- read.delim('https://raw.githubusercontent.com/IowaBiostat/data-sets/main/warming/warming.txt')
y <- df$Raw + runif(nrow(df), -0.005, 0.005)
n <- length(y)

# Trend filtering: order 0 (constant)
# You could also use flsa for this
fit <- trendfilter(y, ord = 0)

# Plot solutions
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, coef(fit, lambda = lam[i])$beta, lwd = 3, col = 'red', type = 's')
  mtext(bquote(lambda == .(lam[i])))
}

# There used to be a nice package called neariso for nearly isotonic regression
# It's no longer on CRAN, so I'll use CVXR instead

lam <- c(0.01, 0.1)
for (i in 1:length(lam)) {
  # Solve problem
  beta <- Variable(n)
  objective <- mean((y - beta)^2) +
    lam[i] * sum(pos(beta[1:(n - 1)] - beta[2:n]))
  problem <- Problem(Minimize(objective))
  result <- solve(problem)

  plot(
    df$Year, y, pch = 19, col = 'gray', xlab = 'Year', ylab = 'LOTI',
    bty = 'n', las = 1
  )
  lines(df$Year, result$getValue(beta), lwd = 3, col = 'red', type = 's')
  mtext(bquote(lambda == .(lam[i])))
}

# Trend filtering: order 1 (linear)
fit <- trendfilter(y)

# Plot solutions
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, coef(fit, lambda = lam[i])$beta, lwd = 3, col = 'red')
  mtext(bquote(lambda == .(lam[i])))
}