import numpy as np
import os
import torch
from utils.myutils import qualify,parse,load_mat
import random
import time
from Denoisers.DRUNet.models.network_unet import UNetRes
from Inpainters.CRUnet.models.Unet import CRUNet
from dapplus import dap_plus_dctnn
from utils.model_utils import init_x
os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'

# set random seed
np.random.seed(5)
random.seed(5)
torch.manual_seed(5)
torch.cuda.manual_seed_all(5)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("device:{}:{}".format(device.type, device.index))
dn_type =1
dn_choice = 1
scale_rate = 1


# load inpainter
pretrain_path = "./Inpainters/CRUnet/checkpoints/CRUNet_color0.1.pth"#
opt = parse('./options/inpaint_color.json', is_train=False)
model_ipt = CRUNet(opt, phase='test')
# 读取配置文件
model_ipt.load_state_dict(torch.load(pretrain_path), strict=True)
model_ipt.eval()
model_ipt = model_ipt.to(device)

# load denoiser
n_channels = 3
pretrain_path_dn = "./Denoisers/DRUNet/checkpoints/drunet_color.pth"
model_dn = UNetRes(in_nc=n_channels + 1, out_nc=n_channels, nc=[64, 128, 256, 512], nb=4, act_mode='R',
            downsample_mode="strideconv", upsample_mode="convtranspose")
model_dn.load_state_dict(torch.load(pretrain_path_dn), strict=True)
model_dn.eval()
for k, v in model_dn.named_parameters():
    v.requires_grad = False
model_dn = model_dn.to(device)



data_choice ='MSI'
data_name = "ds_balloons_ms"
# Test a group of parameters

for sample_rate in [0.03,0.05,0.1]:
    # the decay rate of noise level in the iterations
    if sample_rate == 0.03:
        sigma = 5
        rho = 0.7
        scale_rate = 0.96
    elif sample_rate ==0.05:
        sigma = 10
        rho = 0.5
        scale_rate = 0.94
    elif sample_rate== 0.1:
        sigma =15
        rho = 0.5
        scale_rate = 0.94

    savepath = "./result/"+data_choice+"/"+data_name+"/dapplus_"+data_name+"_sample_rate{}_sigma{}.mat".format(sample_rate,sigma)
    demodatapath = "./data/MSI/balloons_ms/sample_rate{}_sigma{}.mat".format(sample_rate,sigma)
    if not os.path.exists(savepath.rsplit("/",1)[0]):
        os.makedirs(savepath.rsplit("/",1)[0])

    _, datas = load_mat(demodatapath)
    data_gt, a, data_ob = datas["data_gt"], datas["a"], datas["data_ob"]
    a = (a==1)
    psnr_ob,ssim_ob =qualify(data_ob,data_gt,'Observe', final=True)
    [n1, n2, n3] = data_gt.shape
    Ind = np.zeros_like(data_gt)
    Ind[a.reshape(data_gt.shape)] = 1
    save_result = {}

    # init data

    X = np.zeros([n1,n2,n3],dtype=np.float32)
    X = init_x(data_ob,a,model_dn,model_ipt,dn_type,dn_choice,sigma,data_choice)
    save_result["data_gt"] = data_gt
    save_result["init"] = X.copy()
    save_result["Ind"] = Ind.copy()


    print(
        "*****************************************data_name: " + data_name + " | noise_level:{} | sampling_rate:{} | scale_rate:{} | rho:{}*********************************".format(
            sigma, sample_rate, scale_rate,rho))
    parameter = {}
    parameter["dn_type"] = dn_type
    parameter["data_gt"] = data_gt
    parameter["data_ob"] = data_ob
    parameter["X"]  = X.copy()
    parameter["a"] = a.copy()
    parameter["max_iter"] =50
    thresholding = 0.0001
    parameter["stopthres"] = thresholding
    parameter["DEBUG"] =False
    parameter["DEBUG_iter"] = 50
    parameter["change"] =True
    parameter['tnn_type'] = 'dctnn'
    parameter["ro"] = 1.01
    for i_rho1 in [1]:
        for i_rho2 in [1]:
            for i_mu in [1]:
                i_labda = 1
                i_rho3 = i_rho2
                alpha = [1, 1, 1]
                beta = [1, 1, 1]
                filesavepath = savepath
                parameter["lambda"] = i_labda
                parameter["mu"] = i_mu
                parameter["rho"] = rho
                parameter["rho1"] = i_rho1
                parameter["rho2"] = i_rho2
                parameter["rho3"] = i_rho3
                start = time.clock()
                X_f, Y_f, W_f, Z_f, iter_r = dap_plus_dctnn(parameter, model_ipt, model_dn, data_choice, sigma, scale_rate)
                end = time.clock()
                save_dic = {}
                save_dic['w_psnr'], save_dic['w_ssim'] = qualify(W_f, data_gt, 'W_final', final=True)
                save_result["result"] = W_f
                save_result["metric"] = save_dic
                save_result["time"] = end-start
                print("------------------------data_name: " + data_name + " | noise_level:{} | sampling_rate:{} | scale_rate:{}iter:{},time:{}-------------------------".format(sigma,sample_rate,scale_rate,iter_r, end - start))

                img_test = ((Y_f + W_f) / 2)[:, :, 1]


print("over")