# Fits the weighted fused lasso by ADMM where D is the discrete difference operator on a graph
# D is a sparse matrix of class 'dgCMatrix' [package "Matrix"]

fit_graphfusedlasso_cholcache = function(y, lambda, D, chol_factor = NULL, weights=NULL, initial_values = NULL, iter_max = 10000, rel_tol = 1e-4, alpha=1.0, inflate=2, adaptive=FALSE) {
	require(Matrix)
	
	n = length(y)
	m = nrow(D)
	a = 2*lambda # step-size parameter
		
	if(missing(weights)) {
		weights = rep(1, n)
	}
	
	# Check if we need a Cholesky decomp of system involving graph Laplacian
	if(missing(chol_factor)) {
		L = Matrix::crossprod(D)
		chol_factor = Matrix::Cholesky(L + Matrix::Diagonal(n))
	}

	# Initialize primal and dual variables from warm start
	if(missing(initial_values)) {
		x = rep(0, n) # likelihood term
		z = rep(0, n) # slack variable for likelihood
		r = rep(0, m) # penalty term
		s = rep(0, m) # slack variable for penalty
		u_dual = rep(0,n) # scaled dual variable for constraint x = z
		t_dual = rep(0,m) # scaled dual variable for constraint r = s
	} else {
		x = initial_values$x
		z = initial_values$z
		r = initial_values$r
		s = initial_values$s
		t_dual = initial_values$t_dual
		u_dual = initial_values$u_dual
	}
	
	primal_trace = NULL
	dual_trace = NULL
	converged = FALSE
	counter = 0
	while(!converged & counter < iter_max) {
		
		# Update x
		x = {weights * y + a*(z - u_dual)}/{weights + a}
		x_accel = alpha*x + (1-alpha)*z
		
		# Update constraint term r
		arg = s - t_dual
		if(adaptive) {
			local_lambda = 1/{1+(lambda)*abs(arg)}  # Minimax-concave penalty instead?
		} else {
			local_lambda = lambda
		}
		r = softthresh(arg, local_lambda/a)
		r_accel = alpha*r + (1-alpha)*s
		
		# Projection to constraint set
		arg = x_accel + u_dual + Matrix::crossprod(D, r_accel + t_dual)
		z_new = drop(Matrix::solve(chol_factor, arg))
		s_new = as.numeric(D %*% z_new)
		dual_residual_u = a*(z_new - z)
		dual_residual_t = a*(s_new - s)
		z = z_new
		s = s_new
		
		# Dual update
		primal_residual_x = x_accel - z
		primal_residual_r = r_accel - s
		u_dual = u_dual + primal_residual_x
		t_dual = t_dual + primal_residual_r
		
		# Check convergence
		primal_resnorm = sqrt(mean(c(primal_residual_x, primal_residual_r)^2))
		dual_resnorm = sqrt(mean(c(dual_residual_u, dual_residual_t)^2))
		if(dual_resnorm < rel_tol && primal_resnorm < rel_tol) {
			converged=TRUE
		}
		primal_trace = c(primal_trace, primal_resnorm)
		dual_trace = c(dual_trace, dual_resnorm)
		counter = counter+1
		
		# Update step-size parameter based on norm of primal and dual residuals
		if(primal_resnorm > 5*dual_resnorm) {
			a = inflate*a
			u_dual = u_dual/inflate
			t_dual = t_dual/inflate
		} else if(dual_resnorm > 5*primal_resnorm) {
			a = a/inflate
			u_dual = inflate*u_dual
			t_dual = inflate*t_dual
		}
	}
	list(x=x, r=r, z=z, s=s, u_dual=u_dual, t_dual=t_dual,
		primal_trace = primal_trace, dual_trace=dual_trace, counter=counter)
}