
import torch
import torch_dct
import scipy.io as scio
from utils.myutils import maxmin_normal_torch,maxmin_normal_reverse_torch,rgb_padding_torch,maxmin_normal_torch_general,maxmin_nomral_reverse_torch_general
from Inpainters.CRUnet.main_train_gray_completion import inpaint_uninit_torch,inpaint_color_torch,inpaint_color_torch_range,inpaint_color_torch_ho,inpaint_hsv_torch
from Denoisers.DRUNet.main_dpir_denoising import denoising_drunet,denoising_drunet_color,denoising_drunet_torch,denoising_color_torch_range,denoising_rgb_torch,denoising_color_torch_ho,denoising_drunet_hsv
from Inpainters.LBAM.inpainting import inpaint_struct_rgb
def svt(matrix,rho):
    (n1, n2) = matrix.shape
    max12 = max(n1, n2)
    if n1<n2:
        V, S, U_h = torch.linalg.svd(matrix.T,full_matrices=False)
        S = torch.maximum(S - rho, torch.zeros_like(S))
        tol = max12 * torch.finfo(S.dtype).eps
        r = sum(S > tol)
        S = torch.diag(S[:r])
        U = U_h.T
        U = U[:,:r]
        V = V[:,:r]
        V_h = V.T
    else:
        U, S, V_h = torch.linalg.svd(matrix, full_matrices=False)
        S = torch.maximum(S - rho, torch.zeros_like(S))
        tol = max12 * torch.finfo(S.dtype).eps
        r = sum(S > tol)
        S = torch.diag(S[:r])
        U = U[:,:r]
        V_h = V_h[:r, :]

    X = U @ S @ V_h

    return X

def tnn_prox(Y, rho,tnn_type = 'tnn'):
    (n1, n2, n3) = Y.shape
    max12 = max(n1, n2)
    if tnn_type == 'tnn':
        X = torch.zeros(n1, n2, n3, dtype=torch.complex64, device=Y.device)
        Y = torch.fft.fft(Y.clone())
        for i in range(n3):
            U, S, V_h = torch.linalg.svd(Y[:, :, i])
            S = torch.maximum(S - rho, torch.zeros_like(S))
            tol = max12 * torch.finfo(S.dtype).eps
            r = sum(S > tol)
            S = torch.diag(S[:r])
            X[:, :, i] = U[:, 0:r] @ torch.complex(S,torch.zeros_like(S)) @ V_h[0:r, :]
        X = torch.real(torch.fft.ifft(X.clone()))
    elif tnn_type=='dctnn':
        X = torch.zeros(n1, n2, n3, dtype=torch.float32, device=Y.device)
        # Y = torch_dct.dct(Y.clone(), 'ortho')
        Y = torch_dct.dct(Y.clone())
        for i in range(n3):
            X[:,:,i] = svt(Y[:,:,i],rho)
        X = torch_dct.idct(X.clone())
    return X



def update_X_dap_plus(input_tensor,rho,tnn_type):
    return tnn_prox(input_tensor, rho, tnn_type)


def update_W_dap_plus(X,Y,Z,mu,multi1,multi2,multi3,rho1,rho2,rho3,a_r,data_ob):
    x = X.reshape([-1,1])
    y = Y.reshape([-1,1])
    z = Z.reshape([-1,1])
    multi_1 = multi1.reshape([-1, 1])
    multi_2 = multi2.reshape([-1,1])
    multi_3 = multi3.reshape([-1,1])
    a_fla = a_r.reshape([-1,1])
    o = data_ob.clone().reshape([-1,1])
    B = rho1*x+multi_1+ rho2*y+multi_2
    C = (mu*o + rho3*a_fla*z) + multi_3
    W = (1.0/((mu+rho3)*a_fla+rho1+rho2)) * (B+C*a_fla)

    return W.reshape(data_ob.shape)


def update_Y(W,rho2,multi2,model,dn_type=0,data_choice='video',sigma=15):
    multi_2 = multi2.clone()
    beta2 = 1.0/rho2
    mask = torch.ones_like(W) == 1
    tensor_size = [x for x in W.shape]
    n1,n2 = tensor_size[0:2]
    input_tensor, max_t, min_t = maxmin_normal_torch_general(W - (multi_2 * beta2),mask)
    input_tensor = n1n2_padding_general(input_tensor, tensor_size[0], tensor_size[1])

    if data_choice =='RGB':
        result = denoising_rgb_torch(input_tensor, sigma, model)
    elif data_choice == "HSV_ds":
        result = denoising_drunet_hsv(input_tensor,model,sigma)

    elif data_choice=='video_color':
        input_tensor = torch.reshape(input_tensor,[n1,n2,3,-1])
        result =  denoising_color_torch_ho(input_tensor, model, sigma)
        result = torch.reshape(result,[n1,n2,tensor_size[2]])
    else:
        if dn_type:
            result = denoising_color_torch_range(input_tensor, model, dn_type, sigma)
        else:
            result = denoising_drunet_torch(input_tensor, model,sigma)
    result = maxmin_nomral_reverse_torch_general(result,mask,max_t,min_t)
    result=result[:n1,:n2,:]
    return result

def n1n2_padding(tensor,n1,n2):
    if n1%8:
        tensor = rgb_padding_torch(tensor, 1,8-(n1%8))
    if n2%8:
        tensor = rgb_padding_torch(tensor, 2,8-(n2%8))
    return tensor.clone()
def n1n2_padding_general(tensor,n1,n2):
    tensor_size = [x for x in tensor.shape]

    tensor_new = tensor.clone()
    tensor_new = tensor_new.reshape([tensor_size[0],tensor_size[1],-1])
    if n1%8:
        tensor_new = rgb_padding_torch(tensor_new, 1,8-(n1%8))
    if n2%8:
        tensor_new = rgb_padding_torch(tensor_new, 2,8-(n2%8))

    [n1_new,n2_new] = tensor_new.shape[0:2]
    tensor_size[0:2] = [n1_new,n2_new]
    return tensor_new.reshape(tensor_size)


def single_mode_ipt_general(tensor,mask,model,data_chocie='video',ipt_type='admm'):
    tensor_size = tensor.shape
    tensor_ipt, max_t12, min_t12 = maxmin_normal_torch_general(tensor, mask.bool())

    tensor_ipt = n1n2_padding_general(tensor_ipt,tensor_size[0],tensor_size[1])
    mask_ipt =  n1n2_padding_general(mask,tensor_size[0],tensor_size[1])
    n1_new,n2_new = tensor_ipt.shape[0:2]

    if data_chocie=='HSV_ds':
        tensor_ipt = inpaint_hsv_torch(tensor_ipt, mask_ipt, model)
    elif data_chocie =='RGB':
        if ipt_type == 'strc':
            tensor_ipt = inpaint_struct_rgb(tensor_ipt.cuda(),mask_ipt.float().cuda(),model)
    else:
        tensor_ipt = inpaint_uninit_torch(tensor_ipt, mask_ipt, model)

    tensor_ipt = tensor_ipt[:tensor_size[0], :tensor_size[1], :]
    tensor_ipt = maxmin_nomral_reverse_torch_general(tensor_ipt,mask.bool(),max_t12,min_t12)
    return tensor_ipt.clone()

def single_mode_ipt(tensor,mask,model,data_chocie='video',ipt_type='admm'):
    [n1, n2, n3] = tensor.shape
    tensor_ipt, max_t12, min_t12 = maxmin_normal_torch(tensor, mask.bool())
    tensor_ipt = n1n2_padding(tensor_ipt,n1,n2)
    mask_ipt = n1n2_padding(mask,n1,n2)
    if data_chocie=='MSI':
        if ipt_type == 'admm':
            tensor_ipt = inpaint_color_torch_range(tensor_ipt.cuda(), mask_ipt.cuda(), model)
        else:
            tensor_ipt = inpaint_color_torch(tensor_ipt.cuda(), mask_ipt.cuda(), model)
    elif data_chocie =='video_color':
        input_tensor = torch.reshape(tensor_ipt,[n1,n2,3,-1])
        input_mask =  torch.reshape(mask_ipt,[n1,n2,3,-1])
        tensor_ipt = inpaint_color_torch_ho(input_tensor, input_mask, model)
        tensor_ipt = torch.reshape(tensor_ipt,[n1,n2,n3])
    elif data_chocie =='RGB':
        if ipt_type == 'strc':
            tensor_ipt = inpaint_struct_rgb(tensor_ipt.cuda(),mask_ipt.float().cuda(),model)
    else:
        tensor_ipt = inpaint_uninit_torch(tensor_ipt, mask_ipt, model)

    tensor_ipt = tensor_ipt[:n1, :n2, :]
    tensor_ipt = maxmin_normal_reverse_torch(tensor_ipt,mask.bool(),max_t12,min_t12)
    return tensor_ipt.clone()

def update_z(W,multi3,rho3,a,a_r,model,data_choice='video'):
    multi_3=multi3.clone()
    b = torch.zeros_like(W)
    m_3 = torch.zeros_like(W)
    b[a.reshape(W.shape)] = W[a.reshape(W.shape)]
    m_3[a.reshape(W.shape)] = multi_3[a.reshape(W.shape)]
    dividend = 1.0/(rho3)
    b = b - (m_3 * dividend)
    mask = (a_r.reshape(W.shape)).clone()
    if data_choice =='RGB':
        result = single_mode_ipt(b, mask, model, data_choice, 'strc')
    elif data_choice=='HSV_ds':
        result = single_mode_ipt_general(b, mask, model, data_choice)
    else:
        result = single_mode_ipt(b,mask,model,data_choice,'admm')
    result[a.reshape(W.shape)] = b[a.reshape(W.shape)]
    return result


def init_x(data_ob, a, model_dn, model_ipt, dn_type=0,dn_choice=1,sigma = 10,data_choice = 'video',ipt_type= 'init'):
    # test dn2ipt or ipt2dn
    [n1, n2] = data_ob.shape[0:2]
    b = torch.tensor(data_ob.copy()).cuda()
    mask = (a.reshape(data_ob.shape)).copy()
    mask = torch.tensor(mask).cuda()
    b_new = single_mode_ipt(b,mask,model_ipt,data_choice,ipt_type)
    b_new = n1n2_padding_general(b_new,n1,n2)
    if data_choice == 'RGB':
        result_dn = denoising_rgb_torch(b_new,sigma,model_dn)
    else:
        if dn_type:
            result_dn = denoising_drunet_color(b_new, model_dn,dn_type,sigma)
        else:
            result_dn = denoising_drunet(b_new, model_dn,sigma)
    result_dn =result_dn[: n1,:n2,:]
    return result_dn

def init_x_vc(data_ob, a, model_dn, model_ipt, dn_type=0,dn_choice=1,sigma = 10,data_choice = 'video',ipt_type= 'init'):
    [n1, n2] = data_ob.shape[0:2]
    b = torch.tensor(data_ob.copy()).cuda()
    mask = (a.reshape(data_ob.shape)).copy()
    mask = torch.tensor(mask).cuda()
    b_new = single_mode_ipt(b,mask,model_ipt,data_choice,ipt_type)
    result_dn =b_new[: n1,:n2,:]
    return result_dn


def load_mat(path):

    data = scio.loadmat(path)
    return data.keys(),data



