import cv2
import numpy as np
from scipy import signal
import torch
import os


PIL_CONVERT_COLOR = {
    'RGB': lambda pil_image: pil_image.convert("RGB") if pil_image.mode != 'RGB' else pil_image,
    'full_YCbCr': lambda pil_image: pil_image.convert("YCbCr") if pil_image.mode != 'YCbCr'  else pil_image,
    'full_Y': lambda pil_image: pil_image.convert("YCbCr").getchannel(0) if pil_image.mode != 'YCbCr'  else pil_image.getchannel(0),
    'sdtv_Y': lambda pil_image: _rgb2ycbcr(np.array(pil_image))[:,:,0] if pil_image.mode == 'RGB' else NotImplementedError(f"{pil_image.mode} to Y"),
    'sdtv2_Y': lambda pil_image: rgb2y(np.array(pil_image)) if pil_image.mode == 'RGB' else NotImplementedError(f"{pil_image.mode} to Y"),
    'L': lambda pil_image: pil_image.convert("L") if pil_image.mode != 'L'  else pil_image,
}

def _rgb2ycbcr(img, maxVal=255):
    O = np.array([[16],
                  [128],
                  [128]])
    T = np.array([[0.256788235294118, 0.504129411764706, 0.097905882352941],
                  [-0.148223529411765, -0.290992156862745, 0.439215686274510],
                  [0.439215686274510, -0.367788235294118, -0.071427450980392]])

    if maxVal == 1:
        O = O / 255.0

    t = np.reshape(img, (img.shape[0] * img.shape[1], img.shape[2]))
    t = np.dot(t, np.transpose(T))
    t[:, 0] += O[0]
    t[:, 1] += O[1]
    t[:, 2] += O[2]
    ycbcr = np.reshape(t, [img.shape[0], img.shape[1], img.shape[2]])

    return ycbcr


def rgb2y(im):
    """
    this impl:
     0.301,  0.586,  0.113  = 77/256, 150/256, 29/256
    -0.172, -0.340,  0.512  = -44/256, -87/256, 131/256
     0.512, -0.430, -0.082  = 131/256, -110/256, -21/256
    
    ycbcr 601 sdtv spec[1]:
     0.299,  0.587,  0.114
    -0.172, -0.339,  0.511
     0.511, -0.428, -0.083
    
    ycbcr 601 sdtv spec[2]:
     0.299,  0.587,  0.114
    -0.169, -0.331,  0.5
     0.5,   -0.419, -0.081
    
    [1] Video Demystified A Handbook for the Digital Engineer 4th ed - keith Jack, Chapter 3
    [2] Color Space Conversions Adrian Ford, Alan Roberts
    """
    im = im.astype(np.float32)
    R, G, B = im[:,:,0], im[:,:,1], im[:,:,2]
    Y =  77/256*R + 150/256*G +  29/256*B
    # [1] Note that 8-bit YCbCr and R'G'B' data should be saturated a
    return Y.clip(0,255).astype(np.uint8)

def rgb2yuv(im):
    """
    this impl:
     0.301,  0.586,  0.113  = 77/256, 150/256, 29/256
    -0.172, -0.340,  0.512  = -44/256, -87/256, 131/256
     0.512, -0.430, -0.082  = 131/256, -110/256, -21/256
    
    ycbcr 601 sdtv spec[1]:
     0.299,  0.587,  0.114
    -0.172, -0.339,  0.511
     0.511, -0.428, -0.083
    
    ycbcr 601 sdtv spec[2]:
     0.299,  0.587,  0.114
    -0.169, -0.331,  0.5
     0.5,   -0.419, -0.081
    
    [1] Video Demystified A Handbook for the Digital Engineer 4th ed - keith Jack, Chapter 3
    [2] Color Space Conversions Adrian Ford, Alan Roberts
    """
    im = im.astype(np.float32)
    R, G, B = im[:,:,0], im[:,:,1], im[:,:,2]
    Y =  77/256*R + 150/256*G +  29/256*B
    U = -44/256*R -  87/256*G + 131/256*B
    V = 131/256*R - 110/256*G -  21/256*B
    Y, U, V = Y, U + 128, V + 128
    # [1] Note that 8-bit YCbCr and R'G'B' data should be saturated at the 0 and 255 levels to avoid underflow and overflow
    return np.stack([Y,U,V], axis=-1).clip(0,255).astype(np.uint8)

def yuv2rgb(im):
    """
    this impl:
    1,  0,      1.406     = 1, 0, 360/256
    1, -0.344, -0.719     = 1, -88/256, -184/256
    1,  1.777,  0         = 1, 455/256, 0
    
    ycbcr 601 sdtv spec[1]:
    1,      0,  1.371
    1, -0.336, -0.698
    1,  1.732,  0
    
    ycbcr 601 sdtv spec[2]:
    1,      0,  1.403
    1, -0.344, -0.714
    1,  1.773,  0
    
    [1] Video Demystified A Handbook for the Digital Engineer 4th ed - keith Jack, Chapter 3
    [2] Color Space Conversions Adrian Ford, Alan Roberts
    """
    im = im.astype(np.float32)
    Y, Ud, Vd = im[:,:,0], im[:,:,1]-128, im[:,:,2]-128
    R = Y + 360/256*Vd
    G = Y -  88/256*Ud - 184/256*Vd
    B = Y + 455/256*Ud
    # [1] Note that 8-bit YCbCr and R'G'B' data should be saturated at the 0 and 255 levels to avoid underflow and overflow
    return np.stack([R, G, B], axis=-1).clip(0,255).astype(np.uint8)