import numpy as np
import math

def pack_image(image, file, pack_fmt):
    ## txt
    if pack_fmt.lower() in ['txt']:
        image = np.array(image).reshape((-1,1))
        if 'U' in str(image.dtype):
            index = np.where(image=='xx')[0]
            if len(index) > 0:
                image[index] = 0
        np.savetxt(file,image.astype('float'),fmt="%.8f")
    
    ## bin
    elif pack_fmt.lower() in ['bin']:
        image = np.array(image).reshape((-1,1))
        if 'U' in str(image.dtype):
            index = np.where(image=='xx')[0]
            if len(index) > 0:
                image[index] = 0
        image.astype("uint8").tofile(file)

    ## hex
    elif pack_fmt.lower() in ['hex']:
        byte_per_line=16

        if(image.ndim >= 3):
            h, w, c = image.shape
            size = int(h*w)
        elif(image.ndim == 2):
            h, w = image.shape
            size = int(h*w)
            c = 1
        elif(image.ndim == 1):
            size = int(image.shape[0])
            c = 1
        pix_per_line = int(byte_per_line / c)
        output_line = math.ceil((size) / pix_per_line)
        image_f = image.reshape((size, c))

        with open(file, "w") as f:
            for i in range(output_line):
                pixels = ""
                # ... | pixel 3 | pixel 2 | pixel 1 | pixel 0
                for j in range((i+1)*pix_per_line-1, i*pix_per_line-1, -1):
                    # c3 -> byte 3 | c2 -> byte 2 | c1 -> byte 1 | c0 -> byte 0
                    for k in range(c-1, -1, -1):
                        if j >= image_f.shape[0]:
                            pixels = pixels + '00'
                        elif image_f[j, k] == 'xx':
                            pixels = pixels + 'xx'
                        else:
                            pixels = pixels + str_fill(hex((image_f[j, k]).astype('uint8')).lstrip("0x"))
                f.write(pixels + "\n")

def str_fill(value):
    if len(value) == 1:
        value = "0" + value
    elif len(value) == 0:
        value = "00"
    return value

def reorder_image(image, image_fmt, order=None):
    # ============================================================================================
    #                c3    c2    c1    c0
    #   RGBA          A     B     G     R  -> np.array
    #
    #               byte3 byte2 byte1 byte0      
    # order 0:       R     G     B     A
    # order 1:       R     B     G     A
    # order 2:       G     R     B     A
    # order 3:       G     B     R     A
    # order 4:       B     R     G     A
    # order 5:       B     G     R     A  -> IE ARGB Mode
    # order 8:       A     R     G     B
    # order 9:       A     R     B     G
    # order 10:      A     G     R     B
    # order 11:      A     G     B     R
    # order 12:      A     B     R     G
    # order 13:      A     B     G     R  -> IE RGBA Mode
    # ============================================================================================
    if image_fmt.lower() in ['rgba','rgba8888']:
        if order is None:
            order = 13
        assert(image.ndim == 3)
        if image.shape[2]==3:
            image = np.concatenate((image, np.zeros((image.shape[0], image.shape[1], 1), dtype=np.uint8)), axis=2)
        assert(image.shape[2] == 4)
        h,w,c = image.shape
        image_f = image.reshape((h*w,c))
        image_out = []
        for i in range(h*w):
            if order == 0:
                byte0 = image_f[i,3]
                byte1 = image_f[i,2]
                byte2 = image_f[i,1]
                byte3 = image_f[i,0]
            elif order == 1:
                byte0 = image_f[i,3]
                byte1 = image_f[i,1]
                byte2 = image_f[i,2]
                byte3 = image_f[i,0]
            elif order == 2:
                byte0 = image_f[i,3]
                byte1 = image_f[i,2]
                byte2 = image_f[i,0]
                byte3 = image_f[i,1]
            elif order == 3:
                byte0 = image_f[i,3]
                byte1 = image_f[i,0]
                byte2 = image_f[i,2]
                byte3 = image_f[i,1]
            elif order == 4:
                byte0 = image_f[i,3]
                byte1 = image_f[i,1]
                byte2 = image_f[i,0]
                byte3 = image_f[i,2]
            elif order == 5:
                byte0 = image_f[i,3]
                byte1 = image_f[i,0]
                byte2 = image_f[i,1]
                byte3 = image_f[i,2]
            elif order == 8:
                byte0 = image_f[i,2]
                byte1 = image_f[i,1]
                byte2 = image_f[i,0]
                byte3 = image_f[i,3]
            elif order == 9:
                byte0 = image_f[i,1]
                byte1 = image_f[i,2]
                byte2 = image_f[i,0]
                byte3 = image_f[i,3]
            elif order == 10:
                byte0 = image_f[i,2]
                byte1 = image_f[i,0]
                byte2 = image_f[i,1]
                byte3 = image_f[i,3]
            elif order == 11:
                byte0 = image_f[i,0]
                byte1 = image_f[i,2]
                byte2 = image_f[i,1]
                byte3 = image_f[i,3]
            elif order == 12:
                byte0 = image_f[i,1]
                byte1 = image_f[i,0]
                byte2 = image_f[i,2]
                byte3 = image_f[i,3]
            elif order == 13:
                byte0 = image_f[i,0]
                byte1 = image_f[i,1]
                byte2 = image_f[i,2]
                byte3 = image_f[i,3]
            image_out.append(byte0)
            image_out.append(byte1)
            image_out.append(byte2)
            image_out.append(byte3)

        image_out = np.array(image_out,dtype='uint8')
        return image_out

    # ============================================================================================
    #                c3    c2    c1    c0
    #   RGB565        A    B5    G6    R5  -> np.array
    #
    #                    byte3               byte2               byte1               byte0       
    # order 0:     R1[4:0],G1[5:3]     G1[2:0],B1[4:0]	    R0[4:0],G0[5:3]	    G0[2:0],B0[4:0]
    # order 1:     B1[4:0],G1[5:3]	    G1[2:0],R1[4:0]	    B0[4:0],G0[5:3]	    G0[2:0],R0[4:0]  -> IE RGB565 Mode
    # order 2:     R0[4:0],G0[5:3]	    G0[2:0],B0[4:0]	    R1[4:0],G1[5:3]	    G1[2:0],B1[4:0]
    # order 3:     B0[4:0],G0[5:3]	    G0[2:0],R0[4:0]	    B1[4:0],G1[5:3]	    G1[2:0],R1[4:0]
    # order 4:     G1[2:0],B1[4:0]	    R1[4:0],G1[5:3]	    G0[2:0],B0[4:0]	    R0[4:0],G0[5:3]
    # order 5:     G1[2:0],R1[4:0]	    B1[4:0],G1[5:3]	    G0[2:0],R0[4:0]	    B0[4:0],G0[5:3]
    # order 6:     G0[2:0],B0[4:0]	    R0[4:0],G0[5:3]	    G1[2:0],B1[4:0]	    R1[4:0],G1[5:3]
    # order 7:     G0[2:0],R0[4:0]	    B0[4:0],G0[5:3]	    G1[2:0],R1[4:0]	    B1[4:0],G1[5:3]
    # ============================================================================================
    if image_fmt.lower() in ['rgb565']:
        if order is None:
            order = 0
        assert(image.ndim == 3)
        assert((image.shape[2] == 4) | (image.shape[2] == 3))
        h,w,c = image.shape
        image_f = image.reshape((h*w,c))
        image_out = []
        for i in range(0,h*w,2):
            if order == 0:
                byte0 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,2]
                byte1 = ( image_f[i,0] << 3 ) + ( image_f[i,1] >> 3 )
                if i+1 < h*w:
                    byte2 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,2]
                    byte3 = ( image_f[i+1,0] << 3 ) + ( image_f[i+1,1] >> 3 )
                else:
                    byte2 = 'xx'
                    byte3 = 'xx'
            elif order == 1:
                byte0 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,0]
                byte1 = ( image_f[i,2] << 3 ) + ( image_f[i,1] >> 3 )
                if i+1 < h*w:
                    byte2 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,0]
                    byte3 = ( image_f[i+1,2] << 3 ) + ( image_f[i+1,1] >> 3 )
                else:
                    byte2 = 'xx'
                    byte3 = 'xx'
            elif order == 2:
                byte2 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,2]
                byte3 = ( image_f[i,0] << 3 ) + ( image_f[i,1] >> 3 )
                if i+1 < h*w:
                    byte0 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,2]
                    byte1 = ( image_f[i+1,0] << 3 ) + ( image_f[i+1,1] >> 3 )
                else:
                    byte0 = 'xx'
                    byte1 = 'xx'
            elif order == 3:
                byte2 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,0]
                byte3 = ( image_f[i,2] << 3 ) + ( image_f[i,1] >> 3 )
                if i+1 < h*w:
                    byte0 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,0]
                    byte1 = ( image_f[i+1,2] << 3 ) + ( image_f[i+1,1] >> 3 )
                else:
                    byte0 = 'xx'
                    byte1 = 'xx'
            elif order == 4:
                byte0 = ( image_f[i,0] << 3 ) + ( image_f[i,1] >> 3 )
                byte1 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,2]
                if i+1 < h*w:
                    byte2 = ( image_f[i+1,0] << 3 ) + ( image_f[i+1,1] >> 3 )
                    byte3 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,2]
                else:
                    byte2 = 'xx'
                    byte3 = 'xx'
            elif order == 5:
                byte0 = ( image_f[i,2] << 3 ) + ( image_f[i,1] >> 3 )
                byte1 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,0]
                if i+1 < h*w:
                    byte2 = ( image_f[i+1,2] << 3 ) + ( image_f[i+1,1] >> 3 )
                    byte3 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,0]
                else:
                    byte2 = 'xx'
                    byte3 = 'xx'
            elif order == 6:
                byte2 = ( image_f[i,0] << 3 ) + ( image_f[i,1] >> 3 )
                byte3 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,2]
                if i+1 < h*w:
                    byte0 = ( image_f[i+1,0] << 3 ) + ( image_f[i+1,1] >> 3 )
                    byte1 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,2]
                else:
                    byte0 = 'xx'
                    byte1 = 'xx'
            elif order == 7:
                byte2 = ( image_f[i,2] << 3 ) + ( image_f[i,1] >> 3 )
                byte3 = ( (image_f[i,1] & 0x07) << 5 ) + image_f[i,0]
                if i+1 < h*w:
                    byte0 = ( image_f[i+1,2] << 3 ) + ( image_f[i+1,1] >> 3 )
                    byte1 = ( (image_f[i+1,1] & 0x07) << 5 ) + image_f[i+1,0]
                else:
                    byte0 = 'xx'
                    byte1 = 'xx'
            image_out.append(byte0)
            image_out.append(byte1)
            image_out.append(byte2)
            image_out.append(byte3)

        image_out = np.array(image_out)
        return image_out

    # ============================================================================================
    #                c3    c2    c1    c0
    #   YUV           A     V     U     Y  -> np.array
    #   YCBCR         A    Cr    Cb     Y  -> np.array
    # 
    #               byte3 byte2 byte1 byte0      
    # order 0:       Y    U/Cb  V/Cr   A
    # order 1:       Y    V/Cr  U/Cb   A
    # order 2:      U/Cb   Y    V/Cr   A
    # order 3:      U/Cb  V/Cr   Y     A
    # order 4:      V/Cr   Y    U/Cb   A
    # order 5:      V/Cr  U/Cb   Y     A
    # order 8:       A     Y    U/Cb  V/Cr
    # order 9:       A     Y    V/Cr  U/Cb
    # order 10:      A    U/Cb   Y    V/Cr
    # order 11:      A    U/Cb  V/Cr   Y
    # order 12:      A    V/Cr   Y    U/Cb
    # order 13:      A    V/Cr  U/Cb   Y  
    # ============================================================================================
    if image_fmt.lower() in ['ycbcr','ycbcr444','yuv','yuv444'] :
        if order is None:
            order = 13
        assert(image.ndim == 3)
        if image.shape[2]==3:
            image = np.concatenate((image, np.zeros((image.shape[0], image.shape[1], 1), dtype=np.uint8)), axis=2)
        assert(image.shape[2] == 4)

        ## yuv : y~[0,255]; u,v~[-128,128]
        if image_fmt.lower() in ['yuv', 'yuv444'] :
            image[:,:,1:3] += 128

        h,w,c = image.shape
        image_f = image.reshape((h*w,c))
        image_out = []
        for i in range(h*w):
            if order == 0:
                byte0 = image_f[i,3]
                byte1 = image_f[i,2]
                byte2 = image_f[i,1]
                byte3 = image_f[i,0]
            elif order == 1:
                byte0 = image_f[i,3]
                byte1 = image_f[i,1]
                byte2 = image_f[i,2]
                byte3 = image_f[i,0]
            elif order == 2:
                byte0 = image_f[i,3]
                byte1 = image_f[i,2]
                byte2 = image_f[i,0]
                byte3 = image_f[i,1]
            elif order == 3:
                byte0 = image_f[i,3]
                byte1 = image_f[i,0]
                byte2 = image_f[i,2]
                byte3 = image_f[i,1]
            elif order == 4:
                byte0 = image_f[i,3]
                byte1 = image_f[i,1]
                byte2 = image_f[i,0]
                byte3 = image_f[i,2]
            elif order == 5:
                byte0 = image_f[i,3]
                byte1 = image_f[i,0]
                byte2 = image_f[i,1]
                byte3 = image_f[i,2]
            elif order == 8:
                byte0 = image_f[i,2]
                byte1 = image_f[i,1]
                byte2 = image_f[i,0]
                byte3 = image_f[i,3]
            elif order == 9:
                byte0 = image_f[i,1]
                byte1 = image_f[i,2]
                byte2 = image_f[i,0]
                byte3 = image_f[i,3]
            elif order == 10:
                byte0 = image_f[i,2]
                byte1 = image_f[i,0]
                byte2 = image_f[i,1]
                byte3 = image_f[i,3]
            elif order == 11:
                byte0 = image_f[i,0]
                byte1 = image_f[i,2]
                byte2 = image_f[i,1]
                byte3 = image_f[i,3]
            elif order == 12:
                byte0 = image_f[i,1]
                byte1 = image_f[i,0]
                byte2 = image_f[i,2]
                byte3 = image_f[i,3]
            elif order == 13:
                byte0 = image_f[i,0]
                byte1 = image_f[i,1]
                byte2 = image_f[i,2]
                byte3 = image_f[i,3]
            image_out.append(byte0)
            image_out.append(byte1)
            image_out.append(byte2)
            image_out.append(byte3)
            
        image_out = np.array(image_out,dtype='uint8')
        return image_out

    # ============================================================================================
    #   YUV422 & YCBCR422
    #                c3    c2    c1    c0
    #   YUV           A     V     U     Y  -> np.array
    #   YCBCR         A    Cr    Cb     Y  -> np.array
    # 
    #               byte3 byte2 byte1 byte0       
    # order 0:       Y0   U/Cb   Y1   V/Cr
    # order 1:       Y0   V/Cr	  Y1   U/Cb
    # order 2:      U/Cb   Y0	 V/Cr   Y1
    # order 3:      V/Cr   Y0	 U/Cb   Y1
    # order 4:       Y1   U/Cb	  Y0   V/Cr
    # order 5:       Y1   V/Cr	  Y0   U/Cb -> IE UYVY Mode
    # order 6:      U/Cb	Y1	 V/Cr	Y0
    # order 7:      V/Cr	Y1	 U/Cb	Y0  -> IE YUYV Mode
    # ============================================================================================
    if image_fmt.lower() in [ 'ycbcr422','yuv422'] :
        if order is None:
            order = 7
        assert(image.ndim == 3)
        assert((image.shape[2] == 4) | (image.shape[2] == 3))

        ## yuv : y~[0,255]; u,v~[-128,128]
        if image_fmt.lower() in ['yuv422'] :
            image[:,:,1:3] += 128

        h,w,c = image.shape
        image_f = image.reshape((h*w,c))
        image_out = []
        image_planar_Y = []
        image_planar_U = []
        image_planar_V = []
        ##
        for i in range(0,h*w,2):
            if order == 0:
                if i+1 < h*w:
                    byte0 = image_f[i+1,2]
                    byte1 = image_f[i+1,0]
                    byte2 = image_f[i,1]
                    byte3 = image_f[i,0]
                else:
                    byte0 = image_f[i,2]
                    byte1 = 'xx'
                    byte2 = image_f[i,1]
                    byte3 = image_f[i,0]
            elif order == 1:
                if i+1 < h*w:
                    byte0 = image_f[i+1,1]
                    byte1 = image_f[i+1,0]
                    byte2 = image_f[i,2]
                    byte3 = image_f[i,0]
                else:
                    byte0 = image_f[i,1]
                    byte1 = 'xx'
                    byte2 = image_f[i,2]
                    byte3 = image_f[i,0]
            elif order == 2:
                if i+1 < h*w:
                    byte0 = image_f[i+1,0]
                    byte1 = image_f[i+1,2]
                    byte2 = image_f[i,0]
                    byte3 = image_f[i,1]
                else:
                    byte0 = 'xx'
                    byte1 = image_f[i,2]
                    byte2 = image_f[i,0]
                    byte3 = image_f[i,1]
            elif order == 3:
                if i+1 < h*w:
                    byte0 = image_f[i+1,0]
                    byte1 = image_f[i+1,1]
                    byte2 = image_f[i,0]
                    byte3 = image_f[i,2]
                else:
                    byte0 = 'xx'
                    byte1 = image_f[i,1]
                    byte2 = image_f[i,0]
                    byte3 = image_f[i,2]              
            elif order == 4:
                if i+1 < h*w:
                    byte0 = image_f[i,2]
                    byte1 = image_f[i,0]
                    byte2 = image_f[i+1,1]
                    byte3 = image_f[i+1,0]
                else:
                    byte0 = image_f[i,2]
                    byte1 = image_f[i,0]
                    byte2 = image_f[i,1]
                    byte3 = 'xx'              
            elif order == 5:
                if i+1 < h*w:
                    byte0 = image_f[i,1]
                    byte1 = image_f[i,0]
                    byte2 = image_f[i+1,2]
                    byte3 = image_f[i+1,0]
                else:
                    byte0 = image_f[i,1]
                    byte1 = image_f[i,0]
                    byte2 = image_f[i,2]
                    byte3 = 'xx'              
            elif order == 6:
                if i+1 < h*w:
                    byte0 = image_f[i,0]
                    byte1 = image_f[i,2]
                    byte2 = image_f[i+1,0]
                    byte3 = image_f[i+1,1]
                else:
                    byte0 = image_f[i,0]
                    byte1 = image_f[i,2]
                    byte2 = 'xx'
                    byte3 = image_f[i,1]            
            elif order == 7:
                if i+1 < h*w:
                    byte0 = image_f[i,0]
                    byte1 = image_f[i,1]
                    byte2 = image_f[i+1,0]
                    byte3 = image_f[i+1,2]
                else:
                    byte0 = image_f[i,0]
                    byte1 = image_f[i,1]
                    byte2 = 'xx'
                    byte3 = image_f[i,2]
            
            ##
            if order == 'planar':
                if i+1 < h*w:
                    image_planar_Y.append(image_f[i,0])
                    image_planar_Y.append(image_f[i+1,0])
                    image_planar_U.append(image_f[i,1])
                    image_planar_V.append(image_f[i+1,2])
                else:
                    image_planar_Y.append(image_f[i,0])
                    image_planar_U.append(image_f[i,1])
                    image_planar_V.append(image_f[i,2])    
            else:
                image_out.append(byte0)
                image_out.append(byte1)
                image_out.append(byte2)
                image_out.append(byte3)

        ##
        if order == 'planar':
            image_out = np.concatenate([np.array(image_planar_Y),np.array(image_planar_U),np.array(image_planar_V)]) 
        else:
            image_out = np.array(image_out)

        ##
        return image_out
    # ============================================================================================
    #   YUV420 & YCBCR420
    # ============================================================================================
    if image_fmt.lower() in [ 'ycbcr420','yuv420'] :
        assert(image.ndim == 3)
        assert((image.shape[2] == 4) | (image.shape[2] == 3))
        h,w,c = image.shape
        assert(h%2==0)
        assert(w%2==0)
        image_f = image.reshape((h*w,c))
        image_out = []
        image_planar_Y = []
        image_planar_U = []
        image_planar_V = []

        for y in range(h):
            for x in range(w):
                if( (x % 2)==0 ) & ( (y % 2)==0 ):
                    image_planar_Y.append(image[y,x,0])
                    image_planar_U.append(image[y,x,1])
                    image_planar_V.append(image[y,x,2])
                else:
                    image_planar_Y.append(image[y,x,0])

        ##
        if order == 'planar':
            image_out = np.concatenate([np.array(image_planar_Y),np.array(image_planar_U),np.array(image_planar_V)]) 

        ##
        return image_out

    else:
        return image