import time
import torch
from torch import nn, optim 
from utils import * 
import numpy as np 
import matplotlib.pyplot as plt 
from Model import CNT3D
import hdf5storage
import random

'''
Official Code for 3DCNTNN
For more details, refer to 
"Multidimensional nonlinear transform-based tensor representation for 
high-dimensional image reconstruction" by Zhiwei Shi, et al."

'''


def set_seed(seed):
    random.seed(seed) # Python 随机数生成器
    np.random.seed(seed) # NumPy 随机数生成器
    torch.manual_seed(seed) # PyTorch CPU 随机数生成器
    torch.cuda.manual_seed(seed) # 当前 GPU 随机数生成器
    torch.cuda.manual_seed_all(seed) # 多 GPU 随机数生成器
    torch.backends.cudnn.deterministic = True # 确保计算确定性
    torch.backends.cudnn.benchmark = False # 禁止自适应优化

set_seed(4)


dtype = torch.cuda.FloatTensor

max_iter = 4001              # 迭代次数

Data_type = 'Brazil'

Ohsi_file = 'CR_Data/' + Data_type + '/' + Data_type +'.mat'
In_file   = 'CR_Data/' + Data_type + '/' +'.mat'
Out_file   = 'CR_Data/' + Data_type + '/' + 'Output/'

# 读取原始数据, 用于计算 PSNR3D
mat = hdf5storage.loadmat(Ohsi_file)
x_name = 'Xt3'
Data_Ohsi = mat[x_name]
Data_Ohsi = torch.from_numpy(Data_Ohsi).type(dtype).cuda()

mask_name = 'maskt3'
mask      = mat[mask_name]
mask = torch.from_numpy(mask).type(dtype).cuda()

nomask_name = 'nomaskt3'
nomask      = mat[nomask_name]
nomask = torch.from_numpy(nomask).type(dtype).cuda()

pre_name = 'Yt3'
Data_in = mat[pre_name]
Obsert = Data_in
Data_in = torch.from_numpy(Data_in).type(dtype).cuda()


# 输入通道, 输出通道
k3 = 2    
k2 = 1  
k1 = 1  
n4_3 = int(k3 * Data_in.shape[2])
n4_2 = int(k2 * Data_in.shape[1])
n4_1 = int(k1 * Data_in.shape[0])

n3 = Data_in.shape[2]
n2 = Data_in.shape[1]
n1 = Data_in.shape[0]


lambda_ = 0.2 * 10e-7    # 可调参数，建议在0.1*10e-7和1*10e-7之间调节  
alpha1 = 1              
alpha2 = 1                
alpha3 = 1                

model = CNT3D(n1,n4_1, n2,n4_2 ,n3,n4_3).cuda() 


print("3DCNTNN, 区域遗失填充实验")
s = sum([np.prod(list(p.size())) for p in model.parameters()]); 
params = []
Data_in.requires_grad = True
params += [Data_in]
params += [x for x in model.parameters()]   
optimizier = optim.Adam(params, lr=0.003, weight_decay=0)# 10e-9) 
F_norm = nn.MSELoss()

Data_in = Data_in.unsqueeze(0) 
time_sum = 0

for iter in range(max_iter):

    time_start = time.time()       # 记录开始时间

    Data_LR1,Data_Out1, Data_LR2,Data_Out2, Data_LR3,Data_Out3 = model(Data_in)
    Data_LR1 = Data_LR1.squeeze(0)
    Data_LR2 = Data_LR2.squeeze(0)
    Data_LR3 = Data_LR3.squeeze(0) 

    Data_Out1 = Data_Out1.squeeze(0)               
    Data_Out2 = Data_Out2.squeeze(0)
    Data_Out3 =Data_Out3.squeeze(0)

    nuc = 0
    try:
        nuc = lambda_* (alpha1*(torch.norm(Data_LR1[int(iter%n4_1),:,:].cuda(),'nuc')) +\
                        alpha2*(torch.norm(Data_LR2[int(iter%n4_2),:,:].cuda(),'nuc')) +\
                        alpha3*(torch.norm(Data_LR3[int(iter%n4_3),:,:].cuda(),'nuc')))   
        nuc = nuc 
    except:
        nuc = lambda_* (alpha1*(torch.norm(Data_LR1[int(iter+1)%n4_1,:,:].cuda(),'nuc')) +\
                        alpha2*(torch.norm(Data_LR2[int(iter+1)%n4_2,:,:].cuda(),'nuc')) +\
                        alpha3*(torch.norm(Data_LR3[int(iter+1)%n4_3,:,:].cuda(),'nuc')))  
        nuc = nuc
    loss = nuc
    Data_Out = (Data_Out3.permute(2,1,0) *alpha3 + Data_Out2.permute(1,0,2) *alpha2 + Data_Out1 *alpha1) / (alpha1 + alpha2 + alpha3)
    loss += F_norm(Data_Out*mask, Data_Ohsi*mask) 
    Data_Out = Data_Out*nomask + Data_Ohsi*mask

    loss.backward(retain_graph=True)   
    optimizier.step()
    optimizier.zero_grad()   
    time_end = time.time()        
    time_sum = time_sum + (time_end - time_start) 
    
    if iter % 100 == 0:

        print('iteration:',iter)
        Data_Out = Data_Out.cpu().detach().numpy()
        psnr1 = "{:.3f}".format(psnr3d(Data_Ohsi[:,:,0:3], Data_Out[:,:,0:3]))
        psnr2 = "{:.3f}".format(psnr3d(Data_Ohsi[:,:,4:7], Data_Out[:,:,4:7]))
        psnr3 = "{:.3f}".format(psnr3d(Data_Ohsi[:,:,8:11], Data_Out[:,:,8:11]))
        print('node1_PSNR3D = ', psnr1, ', node2_PSNR3D = ', psnr2, ', node3_PSNR3D = ', psnr3)
        print(f"耗时: {time_sum:.4f}秒")  

        T = 1.5 #提高展示图像的亮度
        plt.subplot(121)
        plt.imshow(np.clip(np.stack((Obsert[:,:,2],
                                    Obsert[:,:,1],
                                    Obsert[:,:,0]),2)*T, 0,1))
        plt.axis('off')        
        plt.title('Observed')
        plt.subplot(122)
        plt.imshow(np.clip(np.stack((Data_Out[:,:,2],
                                    Data_Out[:,:,1],
                                    Data_Out[:,:,0]),2)*T ,0,1))
        plt.axis('off') 
        plt.title('Reconstructed')
        plt.show(block=False)
        plt.pause(3)  # 显示3秒
        plt.close()
        
        





