import numpy as np
import scipy.fft

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:
        # print(halfn3)
        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])
        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,:])


    # tnn = tnn/n3;
    X = scipy.fft.idct(X,axis=2)
    return(np.real(X))