import numpy as np
import os
import torch
from utils.myutils import dim_prod,load_mat,qualify
import random
import time

from Denoisers.DRUNet.models.network_unet import UNetRes
from dapplus_rgb_strc import dap_plus_dctnn
from utils.model_utils import init_x
from Inpainters.LBAM.models.LBAMModel import LBAMModel



os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'

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
model_ipt = LBAMModel(4, 3)
pretrained_path = "./Inpainters/LBAM/LBAM_NoBN_Places10Classes.pth"
model_ipt.load_state_dict(torch.load(pretrained_path), strict=True)
model_ipt.eval()
model_ipt = model_ipt.to(device)

n_channels = 3
pretrain_path_dn = "./Denoisers/DRUNet/checkpoints/drunet_color.pth"
print("-----------------------Net: DRUNet-color-------------------------")
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 ='RGB'
data_name = '084'
pic_type =1
sigma =15
pic_loss = True
demodatapath = "./data/RGB/084/"+data_name+"_pic_type{}_sigma{}_.mat".format(pic_type,sigma)
savepath = './result/'+data_choice+'/pic/'+data_name+'_pic_type{}_sigma{}.mat'.format(pic_type,sigma)
if not os.path.exists(savepath.rsplit("/",1)[0]):
    os.makedirs(savepath.rsplit("/",1)[0])

_, datas = load_mat(demodatapath)
data_gt, Ind, data_ob = datas["data_gt"], datas["Ind"], datas["data_ob"]
a = (Ind == 1)
psnr_ob, ssim_ob = qualify(data_ob, data_gt, 'Observe', final=True)
a = a.reshape(dim_prod(data_gt), )

[n1, n2, n3] = data_gt.shape
save_result = {}
X = np.zeros([n1,n2,n3],dtype=np.float32)
start = time.clock()
X = init_x(data_ob,a,model_dn,model_ipt,dn_type,dn_choice,sigma,data_choice='RGB',ipt_type='strc')
psnr_init,ssim_init =qualify(X.cpu().numpy(), data_gt, 'Init', final=True)
save_result["data_gt"] = data_gt
save_result["Ind"] = Ind.copy()

for rho in [0.1]:
    for scale_rate in [0.97]:
        print(
            "*****************************************data_name: " + data_name + " | noise_level:{} | pictype:{} | scale_rate:{} | rho:{}*********************************".format(
                sigma, pic_type, scale_rate,rho))
        parameter = {}
        parameter["dn_type"] = dn_type

        parameter["data_gt"] = data_gt
        parameter["data_ob"] = data_ob
        parameter["X"]  = X
        parameter["a"] = a.copy()
        parameter["max_iter"] =15
        thresholding = 0.0001
        parameter["stopthres"] = thresholding
        parameter["DEBUG"] =False
        parameter["change"] =False
        parameter["tnn_type"] = 'dctnn'
        i_labda = 0.00001
        i_mu =1
        i_rho1 = 1
        i_rho2 = i_rho1
        i_rho3 = i_rho1

        filesavepath = savepath
        parameter["lambda"] = i_labda
        parameter["mu"] = i_mu
        parameter["ro"] = 1.02
        parameter["rho"] = rho
        parameter["rho1"] = i_rho1
        parameter["rho2"] = i_rho2
        parameter["rho3"] = i_rho3
        parameter["DEBUG_iter"] = 1
        start = time.clock()
        X_f, Y_f, W_f, Z_f, iter_r = dap_plus_dctnn(parameter, model_ipt, model_dn, data_choice,dn_choice, sigma, scale_rate)
        end = time.clock()
        save_dic = {}
        save_dic['x_psnr'], save_dic['x_ssim'] = qualify(X_f, data_gt, 'X_final', final=True)
        save_result["result"] = X_f
        save_result["metric"] = save_dic
        save_result["time"] = end-start
        print("------------------------data_name: " + data_name + " | noise_level:{} | pic_type:{} | mu:{} | rho1{} |iter:{},time:{}-------------------------".format(sigma,pic_type, i_mu, i_rho1,iter_r, end - start))


print("over")