import os
import numpy as np
import  torch
from collections import OrderedDict
import json
import random
from scipy import interpolate
import scipy.io as scio
from skimage.metrics import structural_similarity,peak_signal_noise_ratio
import scipy.fft
import cv2



def prox_tnn(Y, rho):
    times = 1
    (n1, n2, n3) = Y.shape
    max12 = max(n1, n2)
    X = np.zeros((n1, n2, n3), dtype=complex)
    Y = np.fft.fft(Y, axis=2)
    tnn = 0
    trank = 0

    U, S, V = np.linalg.svd(Y[:, :, 0])
    S = np.maximum(S - rho * times, 0)
    tol = max12 * np.spacing(S.max())
    r = sum(S > tol)
    S = np.diag(S)
    S = S[:r, :r]
    X[:, :, 0] = np.dot(np.dot(U[:, 0:r], S), V[0:r, :])

    if n3 % 2 == 0:
        halfn3 = n3 // 2
    elif n3 % 2 == 1:
        halfn3 = n3 // 2 + 1
    for i in range(1, halfn3):
        U, S, V = np.linalg.svd(Y[:, :, i])
        S = np.maximum(S - rho * times, 0)
        tol = max12 * np.spacing(S.max())
        r = sum(S > tol)
        S = np.diag(S)
        S = S[:r, :r]
        X[:, :, i] = np.dot(np.dot(U[:, 0:r], S), V[0:r, :])
        X[:, :, n3 - i] = np.conj(X[:, :, i])

    if n3 % 2 == 0:
        U, S, V = np.linalg.svd(Y[:, :, halfn3])
        S = np.maximum(S - rho * times, 0)
        tol = max12 * np.spacing(S.max())
        r = sum(S > tol)
        S = np.diag(S)
        S = S[:r, :r]
        X[:, :, halfn3] = np.dot(np.dot(U[:, 0:r], S), V[0:r, :])

    X = np.fft.ifft(X, axis=2)
    return np.real(X)

def dct_prox(Y, rho):
    times = 1
    (n1, n2, n3) = Y.shape
    max12 = max(n1, n2)
    X = np.zeros((n1, n2, n3),dtype = complex)
    Y = scipy.fft.dct(Y, axis=2)
    tnn = 0
    trank = 0

    for i in range(n3):
        U, S, V = np.linalg.svd(Y[:,:,i])
        # print(i,n3-i)
        S =  np.maximum(S-rho*times, 0)
        tol = max12 * np.spacing(S.max())
        r = sum(S > tol)
        S = np.diag(S)
        S = S[:r,:r]
        X[:,:,i] = np.dot(np.dot(U[:,0:r],S),V[0:r,:])

    X = scipy.fft.idct(X,axis=2)
    return np.real(X)

def dim_prod(T):

    return np.product([x for x in T.shape])

def interpolate2(Mask, tensor, method='linear'):

    [n1,n2,n3] = Mask.shape
    out = np.zeros_like(Mask)
    out2 = np.zeros([n2,n1,n3])
    x= np.arange(0,n1,1)
    y = np.arange(0,n2,1)
    [x,y] = np.meshgrid(x,y)

    for i in range(n3):
        mask = Mask[:,:,i]
        temp = tensor[:,:,i]

        temp2 = temp[mask==1].copy()
        [X,Y] = np.where(mask==1)
        X = X.astype(np.double)
        Y = Y.astype(np.double)
        out2[:,:,i] = interpolate.griddata(np.concatenate([X.reshape(X.shape[0],1),Y.reshape(Y.shape[0],1)],axis=1),temp2,(x,y),method=method).copy()

    out = np.transpose(out2,(1,0,2))
    out = np.nan_to_num(out,nan = 128/255)

    return out

def interpolate_tensor_2d(Mask, tensor,pad_width, method='linear'):
    if pad_width>0:
        tensor_inpaint = np.pad(tensor.copy(),((pad_width,pad_width),(pad_width,pad_width),(pad_width,pad_width)),'symmetric')
        mask_inpaint = np.pad(Mask.copy(),((pad_width,pad_width),(pad_width,pad_width),(pad_width,pad_width)),'symmetric')
    else:
        tensor_inpaint =tensor.copy()
        mask_inpaint = Mask.copy()
        # fisrt_transpose
    tensor_st = np.transpose(tensor_inpaint.copy(),(2,1,0))
    mask_st = np.transpose(mask_inpaint.copy(),(2,1,0))
    # second_transpose
    tensor_nd = np.transpose(tensor_inpaint.copy(),(2,0,1))
    mask_nd = np.transpose(mask_inpaint.copy(),(2,0,1))
    tensor_st = np.transpose(interpolate2(mask_st,tensor_st,method),(2,1,0))
    tensor_nd = np.transpose(interpolate2(mask_nd,tensor_nd,method),(1,2,0))
    inpaint_result = (tensor_st + tensor_nd) * 0.5
    if pad_width>0:
        inpaint_result = inpaint_result[pad_width:-pad_width,pad_width:-pad_width,pad_width:-pad_width]
    return inpaint_result

def interpolate_tensor_3d(Mask, tensor, method='linear'):
    tensor_st = np.transpose(tensor.copy(),(2,1,0))
    mask_st = np.transpose(Mask.copy(),(2,1,0))
    tensor_nd = np.transpose(tensor.copy(),(2,0,1))
    mask_nd = np.transpose(Mask.copy(),(2,0,1))
    tensor_rd = tensor.copy()
    mask_rd= Mask.copy()
    tensor_st = np.transpose(interpolate2(mask_st,tensor_st,method),(2,1,0))
    tensor_nd = np.transpose(interpolate2(mask_nd,tensor_nd,method),(1,2,0))
    tensor_rd = interpolate2(mask_rd, tensor_rd, method='linear')

    return (tensor_st+tensor_nd+tensor_rd) / 3

def load_mat(path):

    data = scio.loadmat(path)
    return data.keys(),data
def save_mat(path,data):
    scio.savemat(path,data)

def psnr(P,T,peakval=255.0):

    mse = np.zeros(P.shape[2])
    for i in range(P.shape[2]):
        mse[i]= np.mean((P[:,:,i] - T[:,:,i]) ** 2)
    mse = np.mean(mse)
    psnr = 10 * np.log10(peakval**2/mse)
    # psnr = 20 * np.log10(peakval/mse)
    return psnr


def relative_error(x,x_pre,avg=1):
    xk_norm = np.linalg.norm(x_pre)
    if xk_norm == 0:
        return  np.linalg.norm(x-x_pre)/avg
    else:
        return np.linalg.norm(x-x_pre)/(xk_norm*avg)

def relative_error_torch(x,x_pre,avg=1):
    xk_norm = torch.linalg.norm(x_pre)
    if xk_norm == 0:
        return  torch.linalg.norm(x-x_pre)/avg
    else:
        return torch.linalg.norm(x-x_pre)/(xk_norm*avg)

def qualify(t,gt,name=None,final=True):
    [n1,n2] = t.shape[0:2]
    if len(t.shape)>3:
        t_m = t.reshape([n1,n2,-1])*255
        gt_m= gt.reshape([n1,n2,-1])*255
    else:
        t_m = t*255
        gt_m = gt*255
    psn =  psnr(t_m,gt_m)
    psnr_m = peak_signal_noise_ratio(gt_m,t_m,data_range=255)
    ssim = structural_similarity( gt_m,t_m, multichannel=True, gaussian_weights=True, sigma=1.5, use_sample_covariance=False, data_range=255)
    print(name+"~~psnr:{}\tssim:{}".format(psn,ssim))
    if final:
        return psnr_m,ssim
def qualify_torch(t,gt,name=None,final=True):
    if len(t.shape) > 3:
        t_m = t.cpu().numpy().reshape(t.shape[0],t.shape[1],-1)
        t_m *=255
        gt_m = gt.cpu().numpy().reshape(t.shape[0],t.shape[1],-1)
        gt_m *=255
    else:
        t_m = t.cpu().numpy()*255
        gt_m = gt.cpu().numpy()*255
    psn =  psnr(t_m,gt_m)
    psnr_m = peak_signal_noise_ratio(gt_m,t_m,data_range=255)
    ssim = structural_similarity( gt_m,t_m, multichannel=True, gaussian_weights=True, sigma=1.5, use_sample_covariance=False, data_range=255)
    # t1.permute([2,0,1]).unsqueeze(dim=0)
    # ssim = pytorch_ssim.ssim(t1.permute([2,0,1]).unsqueeze(dim=0),gt1.permute([2,0,1]).unsqueeze(dim=0))
    print(name+"~~psnr:{}\tssim:{}".format(psn,ssim))
    if final:
        return psnr_m,ssim



def rgb_padding(img,dim,padnum):
    [n1,n2,n3] = img.shape
    if dim ==1:
        img_new = np.pad(img.copy(),((0, padnum), (0,0),(0,0)),'symmetric')
    elif dim ==2:
        img_new = np.pad(img.copy(), ((0, 0), (0, padnum),  (0, 0)), 'symmetric')
    return img_new

def rgb_padding_torch(img,dim,padnum):
    [n1,n2,n3] = img.shape
    if dim ==1:
        img_new = np.pad(img.cpu().numpy().copy(),((0, padnum), (0,0),(0,0)),'symmetric')
    elif dim ==2:
        img_new = np.pad(img.cpu().numpy().copy(), ((0, 0), (0, padnum),  (0, 0)), 'symmetric')
    return torch.tensor(img_new,device=img.device)

def parse(opt_path, is_train=True):

    # ----------------------------------------
    # remove comments starting with '//'
    # ----------------------------------------
    json_str = ''
    with open(opt_path, 'r') as f:
        for line in f:
            line = line.split('//')[0] + '\n'
            json_str += line

    opt = json.loads(json_str, object_pairs_hook=OrderedDict)

    opt['is_train'] = is_train
    opt['opt_path'] = opt_path
    return opt

def tensor_vector_padding(tensor,pad_num):
    [n1,n2,n3] = tensor.shape
    X_c = np.zeros([n1,n2, n3+pad_num])
    X_c[:,:,:(0-pad_num)] = tensor[:,:,:]
    X_c[:,:,(0-pad_num)] = tensor[:,:,(0-pad_num)]
    return X_c.copy()

def maxmin_normal(tensor,mask):
    [n1,n2,n3] = tensor.shape
    tensor_new = tensor.copy()
    max_i = np.zeros([n3,1])
    min_i = np.zeros([n3,1])
    for i in range(n3):
        mask_b = mask[:,:,i]
        b = tensor_new[:, :, i][mask_b]
        max_i[i] = np.max(b)
        min_i[i] = np.min(b)
        b = (b-min_i[i])/(max_i[i]-min_i[i])
        tensor_new[:, :, i][mask_b] = b.copy()

    return tensor_new,max_i,min_i
def maxmin_normal_torch(tensor,mask):
    [n1,n2,n3] = tensor.shape
    tensor_new = tensor.clone()
    max_i = torch.zeros([n3,1],device=tensor.device)
    min_i = max_i.clone()
    for i in range(n3):
        mask_b = mask[:,:,i]
        b = tensor_new[:, :, i][mask_b]
        max_i[i] = torch.max(b)
        min_i[i] = torch.min(b)
        b = (b-min_i[i])/(max_i[i]-min_i[i])
        tensor_new[:, :, i][mask_b] = b.clone()
    return tensor_new,max_i,min_i
def maxmin_normal_torch_general(tensor,mask):
    tensor_size= tensor.shape
    tensor_new = tensor.clone()
    tensor_new = tensor_new.reshape([tensor_size[0],tensor_size[1],-1])
    mask_new = mask.reshape([tensor_size[0],tensor_size[1],-1])
    n3 = tensor_new.shape[2]
    max_i = torch.zeros([n3,1],device=tensor.device)
    min_i = max_i.clone()
    for i in range(n3):
        mask_b = mask_new[:,:,i]
        b = tensor_new[:, :, i][mask_b]
        max_i[i] = torch.max(b)
        min_i[i] = torch.min(b)
        b = (b-min_i[i])/(max_i[i]-min_i[i])
        tensor_new[:, :, i][mask_b] = b.clone()
    return tensor_new.reshape(tensor_size),max_i,min_i

def maxmin_nomral_reverse_torch_general(tensor,mask,max_t,min_t):
    tensor_size = tensor.shape
    tensor_new =tensor.clone()
    tensor_new = tensor_new.reshape([tensor_size[0],tensor_size[1],-1])
    n3 = tensor_new.shape[2]
    for i in range(n3):
        tensor_new[:, :, i] *= (max_t[i]-min_t[i])
        tensor_new[:, :, i] += min_t[i]
    return tensor_new.reshape(tensor_size)

def maxmin_normal_reverse_torch(tensor,mask,max_t,min_t):
    [n1,n2,n3] = tensor.shape
    tensor_new = tensor.clone()
    for i in range(n3):
        tensor_new[:, :, i] *= (max_t[i]-min_t[i])
        tensor_new[:, :, i] += min_t[i]
    return tensor_new
def maxmin_normal_reverse(tensor,mask,max_t,min_t):
    [n1,n2,n3] = tensor.shape
    tensor_new = tensor.copy()
    for i in range(n3):

        tensor_new[:, :, i] *= (max_t[i]-min_t[i])
        tensor_new[:, :, i] += min_t[i]

    return tensor_new



def clip_assemble(tensor,basic_r=128,basic_c=128):
    t = tensor.copy()
    [R,C,_] = t.shape
    r = R-basic_r
    c = C-basic_c
    t_lu = t[:basic_r,:basic_c,:]
    t_lb = t[r:,:basic_c,:]
    t_ru = t[:basic_r,c:,:]
    t_rb = t[r:,c:,:]
    t_new = np.concatenate([t_lu,t_lb,t_ru,t_rb],axis=2)
    return t_new

def clip_assemble_reverse(t,n1,n2,n3,basic_r=128,basic_c=128):
     t_lu,t_lb,t_ru,t_rb = np.split(t,4,axis=2)
     t_re = np.zeros([n1,n2,n3])
     r = n1-basic_r
     c = n2-basic_c
     t_re[:basic_r, :basic_c, :] += t_lu
     t_re[r:, :basic_c, :] += t_lb
     t_re[:basic_r, c:, :] += t_ru
     t_re[r:,c:,:] += t_rb
     t_re[:r,c:-c,:] *= 0.5
     t_re[-r:,c:-c,:] *= 0.5
     t_re[r:-r,:c,:] *= 0.5
     t_re[r:-r,-c:,:] *= 0.5
     t_re[r:-r,c:-c,:] *= 0.25
     return t_re

def msi_clip_assemble(ts,basic_r=128,basic_c=128):
    t = ts.copy()
    n3 = ts.shape[2]
    times_r = int(512/basic_r)
    times_c = int(512/basic_c)
    n3_new = n3*times_r*times_c
    t_new = np.zeros([basic_r,basic_c,n3_new],dtype=np.float32)
    k = 0
    for i in range(times_r):
        for j in range(times_c):
            t_new[:,:,k*n3:(k+1)*n3] = t[i*basic_r:(i+1)* basic_r,j*basic_c:(j+1) * basic_c,:]
            k+=1
    return t_new

def msi_clip_assemble_reverse(ts,n3,basic_r=128,basic_c=128):
    t_re = np.zeros([512,512,n3],dtype=np.float32);
    t = ts.copy()
    times_r = int(512/basic_r)
    times_c = int(512/basic_c)
    k = 0
    for i in range(times_r):
        for j in range(times_c):
            t_re[i * basic_r:(i + 1) * basic_r, j * basic_c:(j + 1) * basic_c, :] = t[:,:,k*n3:(k+1)*n3]
            k+=1
    return t_re

def get_stripe_randommask(n1,n2,n3,max_width,num):
    start_points = random.sample(range(0,n1-max_width),num)
    a = np.float32(np.ones([n1,n2,n3]))
    for start_point in start_points:
        width = random.randint(4,max_width)
        a[:, start_point:start_point+width, :] = 0
    return a.copy()

def get_random_block(n1,n2,n3,max_block_size,num,block_rand):
    start_points = np.zeros([num,2])
    start_points[:,0] = random.sample(range(0,n1-max_block_size-1),num)
    start_points[:,1] = random.sample(range(0,n2-max_block_size-1),num)
    a = np.float32(np.ones([n1,n2,n3]))
    for start_point in start_points:
        if block_rand:
            block_size = random.randint(5,max_block_size)
        else:
            block_size = max_block_size
        w,l = np.uint(start_point[0]), np.uint(start_point[1])
        a[w:w+block_size,l:l+block_size,:] = 0
    return a.copy()

def get_pic_mask(n1,n2,n3,pic_type,mask_path):
    if n1<256:
        path = mask_path+'picmask{}video.mat'.format(pic_type)
    elif n1<512:
        path = mask_path+'picmask{}small.mat'.format(pic_type)
    else:
        path = mask_path+'picmask{}.mat'.format(pic_type)
    keys, data = load_mat(path)
    mask = data['mask']
    mask_tensor = np.zeros([n1,n2,n3])
    for i in range(n3):
        mask_tensor[:,:,i] = mask.copy()
    return mask_tensor.copy()