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

    mu = paramter.get("mu")
    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")
    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
    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):
        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)

        if iter_main>=z_max:
            Y = update_Y(W, rho2, multi2, model_dn, dn_type, data_choice, sigma)
        else:
            Y = X

        if sigma_start <= 10:
            if iter_main < 25:
                sigma *= scale_rate
        else:
            if iter_main < 18:
                sigma *= scale_rate

        if iter_main>=z_max:
            Z =  update_z(W,multi3,rho3,a,a_r,model_ipt,data_choice=data_choice)
        else:
            Z = X
        W = update_W_dap_plus(X, Y, Z, mu, multi1, multi2, multi3, rho1, rho2, rho3, a_r, data_ob)
        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')



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