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=""))
}
options(repr.plot.width=4, repr.plot.height=4)

n <- 256

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

imageplot(f0)

sigma <- .08

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

imageplot(clamp(f))

Jmin <- 4

wav  <- function(f){ perform_wavelet_transf(f, Jmin, +1) }
iwav <- function(fw){ perform_wavelet_transf(fw, Jmin, -1) }

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

plot_wavelet(wav(f), Jmin)

psi <- function (s,T){
    d <- dim(s)
    s0 <- 1 - T**2/vapply(as.vector(abs(s)**2), function(a){max(a,1e-9)}, double(1))
    s0 <- vapply(as.vector(s0), function(a){max(a,0)}, double(1))
    if (max(d,0)!=0){s0 <- array(s0, d)}
    return(s0)
}

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

s <- seq(-3, 3, length=1024)

plot(s, s*psi(s,1), "l", col="blue")
lines(s, s, lty=2, col="red")

theta <- function(x,T){ psi(x, T)*x }
ThreshWav <- function(f,T){ iwav(theta(wav(f), T)) }

T <- 1.5*sigma

options(repr.plot.width=4, repr.plot.height=4)
imageplot(clamp(ThreshWav(f, T)))

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

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

## Insert your code here.

imageplot(clamp(fThresh), paste("SNR = ", round(snr(f0, fThresh),1), "dB"))

w <- 4
n <- 256

source("nt_toolbox/toolbox_wavelet_meshes/meshgrid.R")
grid <- meshgrid_4d(seq(1, n-w+1, w), seq(1, n-w+1, w), seq(0, w-1), seq(0, w-1))
X <- grid$X ; Y <- grid$Y ; dX <- grid$Z ; dY <- grid$S
I <- (X + dX-1) + (Y + dY-1)*n

for (k in 1:(n%/%w)){
    for (l in 1:(n%/%w)){
        I[k,l,,] <- t(I[k,l,,])
    }
}

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

assign <- function(M,I,H){
    M_temp <- as.array(M, dim(I))
    M_temp[I+1] <- H
    return(matrix(M_temp, c(n,n), byrow=FALSE))
}

iblock <- function(H){ assign(matrix(rep(0,n*n), c(n,n)), I, H)}

print( paste("Should be 0:", norm(as.matrix(f) - iblock(block(f)))) )

energy <- function(H){
    H_tmp <- H
    for (i in 1:n%/%w){
        for (j in 1:n%/%w){
            H_tmp[i,j,,] <- sqrt( mean(H_tmp[i,j,,]**2) )
        }
    }
    return(H_tmp)
}

Thresh <- function(H,T){ psi(energy(H), T)*H }
ThreshBlock <- function(x,T){ iblock(Thresh(block(x), T)) }

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

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

## Insert your code here.

T <- 1.25*sigma

options(repr.plot.width=6, repr.plot.height=6)
plot_wavelet(ThreshBlock(wav(f), T), Jmin)

ThreshWav <- function(f, T){ iwav(ThreshBlock(wav(f), T)) }

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

imageplot(clamp(ThreshWav(f, T)))

source("nt_solutions/denoisingwav_4_block/exo3.R")

## Insert your code here.

imageplot(clamp(fBlock), paste("SNR =", round(snr(f0, fBlock),1), "dB"))

wav  <- function(f){ perform_wavelet_transf(f, Jmin, + 1, ti=1) }
iwav <- function(fw){ perform_wavelet_transf(fw, Jmin, -1, ti=1) }

fw <- wav(f)
n <-  256
#dim(fw)[1]

grid <- meshgrid_5d(seq(1, n-w+1, w), seq(1, dim(fw)[1]), seq(1, n-w+1, w), seq(0, w-1), seq(0, w-1))
X <- grid$X ; J <- grid$Y ; Y <- grid$Z ; dX <- grid$S ; dY <- grid$U
I <- (X + dX-1) + (Y + dY-1)*n + (J-1)*n**2
for (k in 1:(n%/%w)){
    for (l in 1:(n%/%w)){
        for (m in 1:dim(fw)[1]){
            I[m,k,l,,] <- t(I[m,k,l,,])
        }
    }
}

block <- function(x){ array(as.vector(aperm(x,c(3,2,1)))[I+1], dim(I)) }


assign <- function(M,I,H){
    M_temp <- as.vector(M)
    M_temp[I+1] <- H

    idx <- 0
    for (i in 1:dim(M)[1]){
        for (j in 1:dim(M)[2]){
            for (k in 1:dim(M)[3]){
                idx <- idx + 1
                M[i,j,k] <- M_temp[ idx ]
            }
        }
    }
    return(M)
}


iblock <- function(H){ assign(array(rep(0,dim(fw)[1]*n*n), c(dim(fw)[1],n,n)), I, H)}

energy <- function(H){
    H_tmp <- H
    for (i in 1:(n%/%w)){
        for (j in 1:(n%/%w)){
            for (k in 1:dim(fw)[1]){
                H_tmp[k,i,j,,] <- sqrt( mean(H_tmp[k,i,j,,]**2) )
            }
        }
    }
    return(H_tmp)
}

Thresh <- function(H,T){ psi(energy(H), T)*H }
ThreshBlock <- function(x,T){ iblock(Thresh(block(x), T)) }

ThreshWav <- function(f, T){ iwav(ThreshBlock(wav(f), T)) }

T <- 1.25*sigma

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

imageplot(clamp(ThreshWav(f, T)))

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

source("nt_solutions/denoisingwav_4_block/exo4.R")

## Insert your code here.

imageplot(clamp(fTI), paste("SNR =", round(snr(f0, fTI),1), "dB"))
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>