# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
######## cude setting######
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
#torch.cuda.set_device(3)
if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')
print("Using device {}".format(device))

###### random_seed ##########
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.deterministic = True
# -

import os
import os.path
import numpy as np
import random
import torch
import cv2
import glob
from PIL import Image
import matplotlib.pyplot as plt
from scipy.fftpack import fft,ifft
from scipy.fftpack import dct,idct
from torchsummary import summary
from torchvision import transforms, datasets
import torch
import torchvision
from os import listdir
from os.path import isfile, join
from torch.utils.data import DataLoader,Dataset
from torch.autograd import Variable
from skimage import metrics
import torch.nn as nn
import torch.optim as optim
import math
import pandas as pd
import pickle
from utils import *
from Network import *
from Dataset import *
import seaborn as sns
from tqdm import tqdm
from scipy import ndimage
import argparse

from utils import *
train_img=read_clean_img('./data/images/train',color_mode='gray',img_size=300)
test_img=read_clean_img('./data/images/test',color_mode='gray',img_size=300)


# +
def data_aug_denoise(img,std_values,mode_list,denoise_net,noise_mode):
    
    #the noise_mode:varying according to the row of image or special noise std.
    #denoise_net:the pre_trained denoise model
    #mode_list: type number of augmentation
    #std_values: for special noise, one special std value.
    
    #output the psnr of denoised images and augmented images themselves
    
    np.random.seed(0)
    img_size=img[0].shape
    (w,h,c)=img[0].shape
    noise_mean=0
    
    pic=[]
    psnr_results=np.zeros((len(img),len(std_values),len(mode_list)))
    
    if denoise_net=='DNCNN':
        net=DnCNN_RL(channels=1, num_of_layers=17)
    elif denoise_net=='MemNet':
        net=MemNet(in_channels=1, channels=20, num_memblock=6, num_resblock=4)
    elif denoise_net=='RIDNet':
        net=RIDNET(in_channels=1)
    elif denoise_net=='DNCNN_SFM':
        net=DnCNN_RL(channels=1, num_of_layers=17)
    elif denoise_net=='MemNet_SFM':
        net=MemNet(in_channels=1, channels=20, num_memblock=6, num_resblock=4)
    elif denoise_net=='RIDNet_SFM':
        net=RIDNET(in_channels=1)
    
    model = nn.DataParallel(net).cuda()
    model.load_state_dict(torch.load(os.path.join("./model",denoise_net,'net.pth' )))
    model.eval()
    
    for noise_idx,noise_std in enumerate(std_values):
        np.random.seed(0)
        for idx in range(img.shape[0]):
            noise=add_white_gaussian_noise(noise_mean,noise_std,img_size,noise_mode)
            noisy_img=img[idx]+noise
            for mode_idx in range(len(mode_list)):
                img_aug=augment_mode_choose(noisy_img,mode_list[mode_idx])
                img_aug=uint2single(img_aug)
                img_aug=single2tensor4(img_aug)
                
                INoisy = Variable(img_aug.cuda())
                INoisy = torch.clamp(INoisy, 0., 1.)
                
                with torch.no_grad():
                    NoiseNetwork = model(INoisy)
                    INetwork = tensor2uint(NoiseNetwork)

                    INetwork = inverse_aug(INetwork,mode_list[mode_idx])

                    pic.append(INetwork)
                    psnr_results[idx][noise_idx][mode_idx]=metrics.peak_signal_noise_ratio(INetwork,img[idx],data_range=255.)
                    
    pic=np.array(pic)
    pic=pic.reshape((len(std_values),img.shape[0],len(mode_list),w,h,c),order='C')           
    return pic,psnr_results

def train_spa_data(model_dir,train_data,test_data,train_img,test_img,noise_std_values,mode_list):
    test_out_results_spatial=[]
    for i in range(len(noise_std_values)):
        train_loader=prepare_dataloader(i,train_data,train_img,'spatial',100,shuffle=True,patch_mode=True,win=50,stride=50)
        test_loader=prepare_dataloader(i,test_data,test_img,'spatial',100,shuffle=False,patch_mode=False,win=50,stride=50)
                
        model_net=Ensemble_2(len(mode_list),'spatial').cuda()
        model_net.apply(weights_init_kaiming)
        criterion=nn.MSELoss()
        criterion.cuda()
        optimizer=optim.Adam(model_net.parameters(),lr=0.001)
           
        print("the train process of spatial %d:"%noise_std_values[i])
        train_loss,train_psnr,train_ssim,test_loss,test_psnr,test_ssim,test_out=\
        train_ensemble(model_dir,noise_std_values[i],train_loader,test_loader,model_net,optimizer,criterion,'spatial',False)
        
        test_out_results_spatial.append(test_out)
        #print("the PSNR of train_data_set at baseline model:",np.mean(baseline_train_psnr,axis=0)[i])
        print("the PSNR of train_data_set after ensemble:",train_psnr)
        #print("the PSNR of test_data_set at baseline model:",np.mean(baseline_test_psnr,axis=0)[i])
        print("the PSNR of test_data_set after ensemble",test_psnr)
        print(test_out.shape)
        print("exam",psnr_ini(test_out,test_img))
    test_out_results_spatial=np.array(test_out_results_spatial)
    return test_out_results_spatial


# -

class Ensemble_2(nn.Module):
    def __init__(self,mode_num,ensemble_method):
        super(Ensemble_2,self).__init__()
        self.resblocks = ResBlock(layers_num=3, mode_num=mode_num, ensemble_method=ensemble_method)
        self.seblock = SE_Block(ch_in=mode_num,ch_out=mode_num*2,ensemble_method=ensemble_method)
        
        self.fusion = nn.Conv2d(in_channels=mode_num*2, out_channels=1,kernel_size=1,bias=False)
                                  
    def forward(self,x):
        
        spatical_attention = x * F.softmax(self.resblocks(x), 1)
        
        channel_attention = self.seblock(x)
        return self.fusion(torch.cat((spatical_attention, channel_attention), 1))


# +
import matplotlib.pyplot as plt
import numpy as np

std_values = [50]
noise_mode = 'normal_noise'
denoise_net = 'DNCNN'
mode_list=[0,1,2,3,4,5,6,7,8,9,10,11,12]

np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
_,baseline_train_psnr,baseline_train_ssim=data_aug_denoise(train_img,opt.noise_std_values,[0],opt.denoise_net,opt.noise_mode)
_,baseline_test_psnr,baseline_train_ssim=data_aug_denoise(train_img,opt.noise_std_values,[0],opt.denoise_net,opt.noise_mode)
train_data,_=data_aug_denoise(train_img,std_values,mode_list, denoise_net, noise_mode)
test_data,test_psnr= data_aug_denoise(test_img, std_values,mode_list, denoise_net, noise_mode)
# -

print(np.mean(test_psnr,axis=0))
baseline_train_psnr=25.12
baseline_test_psnr=24.96

model_dir = os.path.join('saved_models', str(denoise_net))
if not os.path.exists(model_dir):
    os.makedirs(model_dir)

test_out_results_spatial=train_spa_data(model_dir,train_data,test_data,train_img,test_img,[50],mode_list,baseline_train_psnr,baseline_test_psnr)

print(test_out_results_spatial.shape)

print(restored.shape)

restored=restored[0]

plt.figure()
plt.xticks([])
plt.yticks([])
plt.imshow(restored[1][0],cmap='gray')
plt.savefig("RIDnet.png",dpi=60, bbox_inches='tight')

plt.figure()
plt.xticks([])
plt.yticks([])
plt.imshow(test_img[1],cmap='gray')
plt.savefig("GT.png",dpi=60, bbox_inches='tight')

a=np.float32(restored[0][0]/255)-np.float32(test_img[0]/255)

print(a)

# +
import matplotlib.pyplot as plt
import numpy as np

std_values = [20]
noise_mode = 'normal_noise'
denoise_net = 'DNCNN'


np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
restored, _= data_aug_denoise(test_img, std_values, [0,1,2,3,4,5,6,7,8,9,10,11,12], denoise_net, noise_mode)
restored = restored[0]
restored = np.transpose(restored,(1,0,2,3,4))
diff = restored - origin
# print(diff.shape)
print(np.where(diff>0))

origin = test_img

plt.figure()
poss = [[150,150],[75,75],[75,225],[225,75],[225,225]]
# pos = [150,150]
f, axes = plt.subplots(len(poss), 1, figsize=(10,10*3))
nbins =50
img_num = 50
for i, pos in enumerate(poss):
    error = diff[:, img_num, :, :, 0]
    error = error[:, pos[0], pos[1]]
    print(error)
    axes[i].hist(error, bins=nbins)

plt.savefig('error_distribution.pdf', bbox_inches='tight')
plt.show()
# -



# + active=""
# m = torch.nn.Softmax2d()
# input = torch.randn(2, 3, 12, 13)
# output = m(input)
# print(output.shape)
# x = output[1,:,1,1]
# print(x.shape)
# print(x.sum())

# + active=""
# from mpl_toolkits.axes_grid1 import make_axes_locatable
# plt.rcParams.update({'font.size': 15})
# plt.figure()
#
# fig, axes = plt.subplots(1, 3, sharey='row', figsize=(15,4))
# std_values = [25]
# noise_mode = 'normal_noise'
#
# for i, denoise_net in enumerate(['DNCNN', 'MemNet', 'RIDNet']):
#     np.random.seed(0)
#     torch.manual_seed(0)
#     torch.cuda.manual_seed_all(0)
#     restored, _= data_aug_denoise(test_img, std_values, [0,1,2,3,4,5,6,7,8,9,10,11,12], denoise_net, noise_mode)
#     restored = np.swapaxes(restored,1,2)
#     origin = test_img
#     
#     error = np.square((restored - origin))
#     error = np.squeeze(error, axis=0)
#     error = np.squeeze(error, axis=-1)
#     error = error.reshape(error.shape[0], -1)
#     
#     corr = np.corrcoef(error)
#     
#     
#     im = axes[i].imshow(corr)
#     axes[i].set_xticks(range(0,13,1))
#     axes[i].set_xticklabels(range(0,13,1), rotation='vertical')
#     axes[i].set_yticks(range(0,13,1))
#     axes[i].set_title(denoise_net)
#
# plt.subplots_adjust(wspace=0.05)
# fig.colorbar(im, ax=axes.ravel().tolist())
# plt.savefig('correlation.pdf', bbox_inches='tight')
# plt.show()
# -

class Ensemble_2(nn.Module):
    def __init__(self,mode_num,ensemble_method):
        super(Ensemble_2,self).__init__()
        self.resblocks = ResBlock(layers_num=3, mode_num=mode_num, ensemble_method=ensemble_method)
        self.seblock = SE_Block(ch_in=mode_num,ch_out=mode_num*2,ensemble_method=ensemble_method)
        
        self.fusion = nn.Conv2d(in_channels=mode_num*2, out_channels=1,kernel_size=1,bias=False)
                                  
    def forward(self,x):
        
        spatical_attention = x * F.softmax(self.resblocks(x), 1)
        
        channel_attention = self.seblock(x)
        return self.fusion(torch.cat((spatical_attention, channel_attention), 1))


model=Ensemble_2(13,'DCT').cuda()

print('Trainable parameters: ', sum(p.numel() for p in model.parameters() if p.requires_grad))



StarcoderdataJupyterScriptsDedupFiltered