# %%
# #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sept 29, 2026

@author: nellypustelnik
"""

#pip install git+https://github.com/deepinv/deepinv.git#egg=deepinv
import deepinv as dinv
import torch
from torchvision.io import read_image
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(0)  # bruit reproductible

"""
Download an image from an file
------------------------------
"""

image_path = "../images/393035.jpg"
image_file = read_image(image_path)
x_true = image_file.unsqueeze(0).to(torch.float32).to(device) / 255.0
dinv.utils.plot([x_true], titles=["Image originale"])

"""
Add Gaussian noise
------------------
"""
sigma = 0.1
physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise(sigma=sigma))
z = physics(x_true)

#%%
"""
Define Prior : TV
------------------
"""
prior     =  dinv.optim.TVPrior()
perf_psnr = dinv.metric.PSNR()

gamma0 = 0.1  # poids de la régularisation TV dans le prox

x_denoised = prior.prox(z, gamma=gamma0)

psnr_in = perf_psnr(x_true, z).item()
psnr_out = perf_psnr(x_true, x_denoised).item()
print(f"PSNR image bruitée : {psnr_in:.2f} dB")
print(f"PSNR image débruitée : {psnr_out:.2f} dB")

dinv.utils.plot(
    [x_true, z, x_denoised],
    titles=["Originale", f"Bruitée ({psnr_in:.1f} dB)", f"Débruitée TV ({psnr_out:.1f} dB)"],
)


#%%
"""
Define Prior : Wavelet
------------------
"""
prior     =  dinv.optim.WaveletPrior(level=3, wv="db4", p=1)
perf_psnr = dinv.metric.PSNR()

gamma0 = 1.5 # poids de la régularisation TV dans le prox

x_denoised = prior.prox(z, gamma=gamma0)

psnr_in = perf_psnr(x_true, z).item()
psnr_out = perf_psnr(x_true, x_denoised).item()
print(f"PSNR image bruitée : {psnr_in:.2f} dB")
print(f"PSNR image débruitée : {psnr_out:.2f} dB")

dinv.utils.plot(
    [x_true, z, x_denoised],
    titles=["Originale", f"Bruitée ({psnr_in:.1f} dB)", f"Débruitée Wav-l1 ({psnr_out:.1f} dB)"],
)


# %%