#!/usr/bin/env python
# coding: utf-8

# In[31]:


import os
import struct
import numpy as np
import json
import argparse
from PIL import Image


# In[40]:

def main_(args):

    rgb565_file_path = args.input_file
    png_file_path = args.output_file
    bin_format =  args.format

    if not png_file_path.endswith(".png"):
        raise TypeError("The output file should be a png file")

    if bin_format not in ["rgb565", "bgr565", "nir8"]:
        raise TypeError("The format of bin file should be rgb565, bgr565 or nir8")

    row = int(args.height)
    col = int(args.width)


    # In[41]:

    if bin_format != "nir8":


        channel_num = 3
        struct_fmt = '1B' 
        struct_len = struct.calcsize(struct_fmt)
        struct_unpack = struct.Struct(struct_fmt).unpack_from

        pixels = row*col
        len_ = pixels * 3


        rgba565 = []
        with open(rgb565_file_path, "rb") as f:
            while True:
                data = f.read(struct_len)
                if not data: break
                s = struct_unpack(data)
                rgba565.append(s[0])

        rgba565 = rgba565[:pixels*2]

        rgb = [0] * len_

        pixel_num = row*col
        pDesTemp = [0] * len(rgb)

        cnt = 0
        for i in range(0, pixel_num*2, 2):
            temp = rgba565[i]
            temp2 = rgba565[i+1]

            if bin_format == "rgb565":
                #R-5
                pDesTemp[cnt] = ((temp2 >>3) << 3)

                #G-6
                cnt += 1
                pDesTemp[cnt] = ((temp & 0xe0) >> 3) + ((temp2 & 0x07) << 5)

                #B-5
                cnt += 1
                pDesTemp[cnt] = ((temp & 0x1f) << 3)

                cnt += 1

            elif bin_format == "bgr565":
                #R-5
                pDesTemp[cnt] = ((temp & 0x1f) << 3) 

                #G-6
                cnt += 1
                pDesTemp[cnt] = ((temp & 0xe0) >> 3) + ((temp2 & 0x07) << 5)

                #B-5
                cnt += 1
                pDesTemp[cnt] = ((temp2 >>3) << 3)

                cnt += 1


        # print(pDesTemp[:100])

        rgb[:] = pDesTemp


        # In[42]:


        rgb_array = np.zeros((row, col, channel_num))
        for m in range(0, row):
            for n in range(0, col):
                for c in range(0, channel_num):
                    src_index = m * col * channel_num + n * channel_num + c
                    rgb_array[m, n, c] = rgb[src_index] 


        # In[43]:


        im = Image.fromarray(rgb_array.astype('uint8'))
        im.save(png_file_path)
    else:

        channel_num = 3
        struct_fmt = '1B' 
        struct_len = struct.calcsize(struct_fmt)
        struct_unpack = struct.Struct(struct_fmt).unpack_from

        pixels = row*col
        len_ = pixels * 3


        nir8 = []
        with open(rgb565_file_path, "rb") as f:
            while True:
                data = f.read(struct_len)
                if not data: break
                s = struct_unpack(data)
                nir8.append(s[0])

        nir8 = nir8[:pixels]

        nir = [0] * len_

        pixel_num = row*col
        pDesTemp = [0] * len(nir)

        cnt = 0
        for i in range(0, pixel_num, 1):
            temp = nir8[i]

            if bin_format == "nir8":
                
                pDesTemp[cnt] = temp

                #G-6
                cnt += 1
                pDesTemp[cnt] = temp
                #B-5
                cnt += 1
                pDesTemp[cnt] = temp

                cnt += 1

            


        # print(pDesTemp[:100])

        nir[:] = pDesTemp

        # print(nir)


        # In[42]:


        nir_array = np.zeros((row, col, channel_num))
        for m in range(0, row):
            for n in range(0, col):
                for c in range(0, channel_num):
                    src_index = m * col * channel_num + n * channel_num + c
                    nir_array[m, n, c] = nir[src_index] 

        im = Image.fromarray(nir_array.astype('uint8'))
        im.save(png_file_path)


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(
        description="script converting rgb565 binary to png file"
        )

    argparser.add_argument(
        '-i',
        '--input_file',
        help="input file should be a bin file"
        )

    argparser.add_argument(
        '-f',
        '--format',
        help="format of the bin file, i.e. rgb565, bgr565"
        )
    
    argparser.add_argument(
        '-o',
        '--output_file',
        help="input file should be a png file"
    )

    argparser.add_argument(
        '-he',
        '--height',
        help="the height of the rgb565"
    )

    argparser.add_argument(
        '-w',
        '--width',
        help="the width of the rgb565"
    )

    args = argparser.parse_args()

    main_(args)