# npls_sps.jl

using SparseArrays: sparse
using LinearAlgebra: I, norm
using MIRTjim: jim

"""
    xx = npls_sps(yy ; niter, beta, delta)
Nonquadratic penalized least-squares denoising of an image `y`
using a "separable paraboloidal surrogates" (SPS) algorithm
* `yy` image to be denoised
* `niter` # of iterations (default 80)
* `beta` regularization parameter (default 2^4)
* `delta` potential function parameter (default 0.5)
"""
function npls_sps(yy ; niter::Int=80, beta::Real=2^4, delta::Real=0.5)

	(nx, ny) = size(yy)
	C = buildC(nx,ny) # create penalty matrix for 1st-order differences

	# \omega "curvature" function for Fair (aka Lange3) potential function
	wt = (t) -> 1 / (1 + abs(t) / delta)
	#wt = (t) -> 1 # for quadratic potential function

	denom = 1 .+ beta * abs.(C)' * (abs.(C) * ones(nx*ny,1)) # SQS version
	# denom = 1 + 8 * beta; % simple version based on matrix 1-norm of C'C

	xx = vec(yy) # initial guess: the noisy image - in a vector

	for ii=1:niter
		Cx = C * xx
		grad = xx - vec(yy) + beta * (C' * (wt.(Cx) .* Cx))
		xx = xx - grad ./ denom # diagonally preconditioned GD
	end
	xx = reshape(xx, size(yy)) # reshape vector back into an image
	return xx
end


"""
    C1 = buildC1(N)
Build a sparse matrix that computes 1D first-order differences
of a 1D signal of length `N`.
"""
function buildC1(N)
	i = 1:(N-1)
	j = 1:(N-1) # row and column indices
	i = [i; i]
	j = [j; j .+ 1]
	s = [ones(N-1); -ones(N-1)] # 1 and -1 values
	C1 = sparse(i, j, s) # matrix rows are [0 ... 0 1 -1 0 ... 0]
	return C1
end


"""
    C = buildC(nx, ny)
Build a sparse matrix that computes first-order differences
between horizontal and vertical neighboring pixels,
for (vectorized) 2D image of size `nx × ny`.
"""
function buildC(nx, ny)
	Cx = buildC1(nx) # [nx-1 × nx]
	Cy = buildC1(ny) # [ny-1 × ny]
	# make it apply to each row of image (respectively each column)
	# and combine horizontal and vertical penalties:
	C = [kron(I(ny), Cx); kron(Cy, I(nx))]
	return C
end


# test denoising
function npls_sps_test()
	xtrue = zeros(60, 40); xtrue[10:20,5:10] .= 30 # basic test image
	yy = xtrue + 5 * randn(size(xtrue)) # add noise
	xhat = npls_sps(yy ; niter=90)
	nrmse = xh -> round(norm(xh - xtrue) / norm(xtrue) * 100, digits=1)
	@show nrmse(yy), nrmse(xhat) # big improvement!
	jim(cat(dims=3, xtrue, yy, xhat))
end

# npls_sps_test()