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


library(imager)
library(png)


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=""))
}

source("nt_toolbox/toolbox_wavelet_meshes/meshgrid.R")

options(repr.plot.width=3.5, repr.plot.height=3.5)

n <- 512

f0 <- load_image("nt_toolbox/data/lena.png", n)

n <- 128

c <- c(100,200)
f0 <- rescale(f0[(c[1]-n%/%2 + 1):(c[1]+n%/%2), (c[2]-n%/%2 + 1):(c[2]+n%/%2),,])

imageplot(f0)

sigma <- 0.04

f <- as.cimg(f0) + sigma*as.cimg(rnorm(n**2))

imageplot(clamp(f))

w <- 3
w1 <- 2*w + 1

grid <- meshgrid_4d(1:n, 1:n, (-w):w, (-w):w)
X <- grid$X ; Y <- grid$Y ; dX <- grid$Z ; dY <- grid$S
X <- X + dX
Y <- Y + dY

X[X < 1] <- 2-X[X < 1] 
Y[Y < 1] <- 2-Y[Y < 1]
X[X > n] <- 2*n-X[X > n]
Y[Y > n] <- 2*n-Y[Y > n]

I <- (X-1) + (Y-1)*n
for (i in 1:(n%/%w)){
    for (j in 1:(n%/%w)){
        I[i,j,,] <- t(I[i,j,,])
    }
}

patch <- function(f){ array(as.vector(f)[I+1], dim(I)) }

P <- patch(f)

options(repr.plot.width=7, repr.plot.height=7)

for (i in 1:16){
    x <- sample(1:n, 1)
    y <- sample(1:n, 1)
    imageplot(P[x,y,,], "", c(4, 4, i))
}

d <- 25

resh <- function(P){ t(array(P, c(n*n,w1*w1))) }

remove_mean <- function(Q){ Q - array(rep(apply(Q, 2, mean), each=w1*w1), c(w1*w1, n*n)) }

P1 <- remove_mean(resh(P))
C <- P1 %*% t(P1)

eg <- eigen(C)
D <- eg$values ; V <- eg$vectors
D <- D[order(-D)]
I <- order(-D)
V <- V[I,]

options(repr.plot.width=3.5, repr.plot.height=3.5)

plot(1:length(D), D, "l", col="blue")

options(repr.plot.width=7, repr.plot.height=7)

for (i in 1:16){
    imageplot(abs(array(V[,i], c(w1,w1))), "", c(4,4,i))
}

iresh <- function(Q){ array( t(Q), c(n,n,d) ) }
descriptor <- function(f){ iresh( t(V[, 1:d]) %*% remove_mean(resh(P)) ) }

H <- descriptor(f)

distance <- function(i){ apply((H - array( rep(H[i[1],i[2],], each=n*n), dim(H) ))**2, c(1,2), sum)/(w1*w1) }

normalize <- function(K){ K/sum(K) }
kernel <- function(i,tau){ normalize(exp(-distance(i)/(2*tau**2))) }

tau <- 0.05
i <- c(84,73)
D <- distance(i)
K <- kernel(i, tau)

options(repr.plot.width=7, repr.plot.height=3.5)

imageplot(D, 'D', c(1, 2, 1))
imageplot(K, 'K', c(1, 2, 2))

q <- 14

selection <- function(i){
    a <- clamp((i[1]-q):(i[1]+q), 0, n-1)
    b <- clamp((i[2]-q):(i[2]+q), 0, n-1)
    return( t(array(c(a,b), c(length(a), 2))) )
}

distance_0 <- function(i, sel){
    H1 <- H[sel[1,]+1,,]
    H2 <- H1[,sel[2,]+1,]
    return(apply((H2 - array( rep(H[i[1]+1,i[2]+1,], each=length(sel[1,])*length(sel[2,])), dim(H2) ))**2, c(1,2), sum)/w1*w1)
}

distance <- function(i){ distance_0(i, selection(i)) }
kernel <- function(i, tau){ normalize(exp(-distance(i)/ (2*tau**2))) }

D <- distance(i)
K <- kernel(i, tau)

options(repr.plot.width=7, repr.plot.height=3.5)

imageplot(D, 'D', c(1, 2, 1))
imageplot(K, 'K', c(1, 2, 2))

NLval_0 <- function(K,sel){
    f_temp <- f[sel[1,]+1,,,]
    return( sum(K*f_temp[, sel[2,]+1]) )
}

NLval <- function(i, tau){ 
    sel <- selection(i)
    K <- kernel(i, tau)
    return(NLval_0(K, sel)) }

grid <- meshgrid_2d(0:(n-1), 0:(n-1))
Y <- grid$X ; X <- grid$Y

arrayfun <- function(f,X,Y){
    n <- dim(X)[1]
    p <- dim(Y)[1]
    R <- matrix(rep(0,n*p), c(n,p))
    for (k in 0:(n-1)){
        for (l in 0:(p-1)){
            R[k+1,l+1] <- f(k, l)
        }
    }
    return(R) }

NLmeans <- function(tau){ arrayfun(function(i1, i2){NLval(c(i1,i2), tau)}, X, Y) }

tau <- 0.03

options(repr.plot.width=3.5, repr.plot.height=3.5)
imageplot(NLmeans(tau))

options(repr.plot.width=7, repr.plot.height=7)

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

## Insert your code here.

options(repr.plot.width=3.5, repr.plot.height=3.5)

imageplot(clamp(fNL))

options(repr.plot.width=7, repr.plot.height=7)

source("nt_solutions/denoisingadv_6_nl_means/exo2.R")

## Insert your code here.
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>