library(pracma)
library(imager)

options(warn=-1) # turns off warnings, to turn on: "options(warn=0)"

# Importing the libraries
for (f in list.files(path="nt_toolbox/toolbox_general/", pattern="*.R")) {
    source(paste("nt_toolbox/toolbox_general/", f, sep=""))
}
for (f in list.files(path="nt_toolbox/toolbox_signal/", pattern="*.R")) {
    source(paste("nt_toolbox/toolbox_signal/", f, sep=""))
}

# Renaming the imported div_2 and grad_2 function
div = div_2
grad = grad_3

As = function(u){-div(u)}
A = function(u){grad(u)}

norm12 = function(u){sum(sum(sqrt((u^2)[,,1] + (u^2)[,,2])))}
J = function(x){norm12(A(x))}

n = 256

name = "nt_toolbox/data/hibiscus.png"
x0 = load_image(name, n)
x0 = resize(x0, size_x=n, size_y=n)
# Two dimensional
x0 = x0[,,1,1]

options(repr.plot.width=5, repr.plot.height=5)
imageplot(clamp(x0))

sigma = 0.1
set.seed(1)
y = x0 + randn(n,n) * sigma

imageplot(clamp(y))

lambda = 0.2

f = function(x){0.5 * norm(x-y)^2}
g = function(x){lambda * J(x)}
E = function(x){f(x) + g(x)}

F = function(u){0.5 * norm(y - As(u))^2 - 0.5 * norm(y)^2}

nablaF = function(u){A(As(u) - y)}

d = function(u)
{
    n = dim(u)[1]
    p = dim(u)[2]

    out = array(0, c(n, p, 2))
    
    for (i in 1:n)
    {
        for (j in 1:p)
        {
            out[i, j,] = norm(u[i, j,])
        }
    }
    return (out)
}

proxG = function(u){u / pmax(d(u) / lambda, 1)}

gamma = 1./5

u = array(0, c(n, n, 2))

u = proxG(u - gamma * nablaF(u))

x = y - As(u)

source("nt_solutions/optim_7_duality/exo1.R")

# Insert your code here.

imageplot(clamp(x))