from scipy.fftpack import dct,idct
import numpy as np

def dct3(gop):
    '''
    gop shape is bath*h*w
    '''
    assert type(gop) == np.ndarray,"輸入必須為numpy對(duì)象"
    
    gop_dct = []
    # 先對(duì)圖像做二維的dct變換
    for frame in gop:
        frame_dct2 = dct(dct(frame.astype(np.double), norm='ortho').T,norm='ortho').T
        gop_dct.append(frame_dct2)
    gop_dct = np.array(gop_dct)
    result = np.zeros_like(gop_dct)
    _,h,w = gop.shape
    
    # 對(duì)幀間做dct變換
    for i in range(h):
        for j in range(w):
            result[:,i,j] = dct(gop_dct[:,i,j],norm='ortho')
    return result
 
    
def idct3(gop_dct):
    '''
    gop_dct shape is bath*h*w
    '''
    assert type(gop_dct) == np.ndarray,"輸入必須為numpy對(duì)象"
    
    # 先對(duì)幀間做idct
    _,h,w = gop_dct.shape
    frame_rec = np.zeros_like(gop_dct)
    
    for i in range(h):
        for j in range(w):
            frame_rec[:,i,j] = idct(gop_dct[:,i,j],norm='ortho')
      
    result = []
    for frame in frame_rec:
        result.append(idct(idct(frame, norm='ortho').T,norm='ortho').T)
    return np.array(result)