
import torch
from utils.model_utils import update_Y,update_z,update_X_dap_plus,update_W_dap_plus
from utils.myutils import relative_error_torch,qualify_torch
def dap_plus_dctnn(paramter,model_ipt,model_dn,data_choice,sigma=15,scale_rate=0.99):

    # 获得相关参数
    # alpha = paramter.get("alpha")
    mu = paramter.get("mu")
    labda = paramter.get("lambda")
    rho = paramter.get("rho")
    rho1 = paramter.get("rho1")
    rho2 = paramter.get("rho2")
    rho3 = paramter.get("rho3")
    max_iter = paramter.get("max_iter")
    DEBUG = paramter.get("DEBUG")
    stopthres = paramter["stopthres"]
    dn_type = paramter["dn_type"]
    change = paramter["change"]
    tnn_type = paramter["tnn_type"]
    ro= paramter.get("ro")
    debug_iter = paramter.get("DEBUG_iter")
    if debug_iter is None:
        debug_iter=10
    # 获得所有运算变量
    a = paramter["a"]
    a = torch.tensor(a).cuda()
    a_r = a.float()
    data_gt = torch.tensor(paramter["data_gt"]).cuda()
    data_ob = torch.tensor(paramter["data_ob"]).cuda()
    X = paramter.get("X")
    if not torch.is_tensor(X):
        X = torch.tensor(paramter.get("X")).cuda()

    W = X.clone()
    Y = X.clone()
    Z = X.clone()
    tensor_size = X.shape

    multi1 = torch.zeros_like(X)
    multi2 = torch.zeros_like(X)
    multi3 = torch.zeros_like(X)

    sigma_start = sigma
    stop_flag = 0
    if data_choice == 'MSI_ds':
        z_max = 0
    elif data_choice =='HSV_ds':
        z_max=20
    elif data_choice=='video':
        z_max=0
    else:
        z_max = 20
    for iter_main in range(max_iter):


        #留个副本
        Xk = X.clone()
        Wk = W.clone()
        Yk = Y.clone()
        Zk = Z.clone()
        # Multipiers_k1k2k = Multipiers_k1k2.clone()
        # multi2k = multi2.clone()
        # multi3k = multi3.clone()

        # start_x = time.clock()
        # updatE_X
        beta1 = 1.0 / rho1
        input_updatex = W - (multi1 * beta1)
        if data_choice=='HSV_ds':
            input_updatex = input_updatex.reshape([tensor_size[0],tensor_size[1],-1])
            X = torch.reshape(update_X_dap_plus(input_updatex,rho,tnn_type),tensor_size)
        else:
            X = update_X_dap_plus(input_updatex,rho,tnn_type)


        # end_x= time.clock()
        # update_Y
        # start_y = time.clock()
        if iter_main>=z_max:#20
            Y = update_Y(W, rho2, multi2, model_dn, dn_type, data_choice, sigma)
        else:
            Y = X
        if iter_main>=max_iter-2:
            Y = update_Y(Y, rho2, torch.zeros_like(X), model_dn, dn_type, "video_color", sigma)
        # end_y = time.clock()

        if sigma_start <= 10:
            if iter_main < 20:
                # if sigma > 4.6:
                sigma *= scale_rate
        else:
            if iter_main < 20:
                sigma *= scale_rate

        if iter_main>=z_max:
            Z =  update_z(W,multi3,rho3,a,a_r,model_ipt,data_choice='video_color')
            # end_z = time.clock()
        else:
            Z = X
        W = update_W_dap_plus(X, Y, Z, mu, multi1, multi2, multi3, rho1, rho2, rho3, a_r, data_ob)
        # end_w = time.clock()


        multi1 += rho1*(X-W)
        multi2 += (rho2*(Y-W))
        multi3 += (rho3*((a_r.reshape(-1,1) * Z.reshape(-1,1)) - (a_r.reshape(-1,1) * W.reshape(-1,1)))).reshape(data_gt.shape)

        if change:

            rho/=ro
            mu *= ro
            rho1 *= ro
            rho2 *= ro
            rho3 *= ro

        #
        if DEBUG and (iter_main+1)%debug_iter==0:
            print('iter-{}'.format(iter_main+1).center(80, '-'))
            qualify_torch(X, data_gt, 'X')
            qualify_torch(Y, data_gt, 'Y')
            qualify_torch(W, data_gt, 'W')
            qualify_torch(Z, data_gt, 'Z')
        #     qualify_torch((X_ws + Z) / 2, data_gt, 'MEAN')
        #     qualify_torch(( X_ws+Y +Z) / 3, data_gt, 'xyz')

        stop_list = [relative_error_torch(X, Xk), relative_error_torch(Y,Yk), relative_error_torch(Z,Zk)]

        stopx = max(stop_list)
        if stopx < 0.001 and stop_flag==0:
            print("1e-3 stop!")
            stop_flag =1
        if stopx < stopthres:
            # X_ws = torch.sum((beta*X_k1k2),dim=0).float()
            # qualify_torch(X_ws, data_gt, 'X_ws')
            # qualify_torch(Y, data_gt, 'Y')
            # qualify_torch(Z, data_gt, 'Z')
            # qualify_torch((X_ws+Y+W+Z)/4, data_gt,'MEAN')
            # qualify_torch((X + Y + W + Z) / 4, data_gt, 'MEAN_init')
            # qualify_torch((X + Z) / 2, data_gt, 'MEAN_XZ')
            # qualify_torch((Z + Y) / 2, data_gt, 'MEAN_ZY')
            break
    # 需要考虑的问题，是否用去噪后的observe部分替代结果变量对应的observe部分。

    return X.cpu().numpy(), Y.cpu().numpy(), W.cpu().numpy(), Z.cpu().numpy(), iter_main
