#!/usr/bin/env python3
from PIL import Image
from scipy import fftpack
import os
import huffmanEncode
import numpy as np
from numpy import *
import sys
import time

DEBUG = False

jpeg_header = 'FFD8FFE000104A46494600010100000100010000' # len = 40
lumQuantTbl_header = 'FFDB004300'
lumQuantTbl_header_dynamic = 'FFDB004310'
chromQuantTbl_header = 'FFDB004301'
chromQuantTbl_header_dynamic = 'FFDB004311'
heightwidth_header = 'FFC0001108'
subsamp = '03011100021101031101'
huffmanTable = 'FFC401A20000000701010101010000000000000000040503020601000708090A0B0100020203010101010100000000000000010002030405060708090A0B1000020103030204020607030402060273010203110400052112314151061361227181143291A10715B14223C152D1E1331662F0247282F12543345392A2B26373C235442793A3B33617546474C3D2E2082683090A181984944546A4B456D355281AF2E3F3C4D4E4F465758595A5B5C5D5E5F566768696A6B6C6D6E6F637475767778797A7B7C7D7E7F738485868788898A8B8C8D8E8F82939495969798999A9B9C9D9E9F92A3A4A5A6A7A8A9AAABACADAEAFA110002020102030505040506040803036D0100021103042112314105511361220671819132A1B1F014C1D1E1234215526272F1332434438216925325A263B2C20773D235E2448317549308090A18192636451A2764745537F2A3B3C32829D3E3F38494A4B4C4D4E4F465758595A5B5C5D5E5F5465666768696A6B6C6D6E6F6475767778797A7B7C7D7E7F738485868788898A8B8C8D8E8F839495969798999A9B9C9D9E9F92A3A4A5A6A7A8A9AAABACADAEAFA'
data_header = 'FFDA000C03010002110311003F00'

zigzagOrder = np.array([0,1,8,16,9,2,3,10,17,24,32,25,18,11,4,5,12,19,26,33,40,48,41,34,27,20,13,6,7,14,21,28,35,42,
                           49,56,57,50,43,36,29,22,15,23,30,37,44,51,58,59,52,45,38,31,39,46,53,60,61,54,47,55,62,63])

def loadDetection(txtFileName):
    txt = open(txtFileName, 'r')
    lines = txt.readlines()
    for i, line in enumerate(lines):
        lines[i] = list(map(int, line.strip().split(' ')[-4:]))
    return lines

def isAround(bounding_box, lefttop):
    x, y = lefttop[0], lefttop[1]

    for left, top, right, bottom in bounding_box:
        if (left - 8 < x) and (x < right + 8) and (bottom - 8 > y) and (y > top + 8) :
            return True
    return False

def getBinNum(frameStr):
    return int(frameStr[5:-4])

def binCompare(fileNameA, fileNameB):
    aNum = getBinNum(fileNameA)
    bNum = getBinNum(fileNameB)
    return aNum - bNum

def main():

    srcFileName = sys.argv[1]
    outFileName = sys.argv[2]
    prevDetections = ''
    if(len(sys.argv)>3):
        prevDetections = sys.argv[3]
        bb = loadDetection(prevDetections)
        print("{} detections loaded.".format(len(bb)))

    t0 = time.time()

    print("Decoding {}...".format(srcFileName))

    f = open(srcFileName,'rb')

    # JPEG header
    target = jpeg_header
    b = f.read(len(target)//2)
    if b == huffmanEncode.hexToBytes(target):
        if DEBUG:
            print("JPEG header good")
    else:
        if DEBUG:
            print("JPEG header bad")
        sys.exit()

    # lumQuantTbl header & value
    target = lumQuantTbl_header
    b = f.read(len(target)//2)
    if b == huffmanEncode.hexToBytes(target):
        if DEBUG:
            print("lumQuantTbl header good. non dynamic")
        b = f.read(64)
        luminanceQuantTbl = np.array(list(b)).reshape([8,8])
    elif b == huffmanEncode.hexToBytes(lumQuantTbl_header_dynamic):
        if DEBUG:
            print("lumQuantTbl header good. dynamic")
        b = f.read(64)
        luminanceQuantTbl = np.array(list(b)).reshape([8,8])
        b = f.read(64)
        luminanceQuantTbl_lo = np.array(list(b)).reshape([8,8])
    else:
        if DEBUG:
            print("lumQuantTbl header bad")
        sys.exit()

    # chromQuantTbl header & value
    target = chromQuantTbl_header
    b = f.read(len(target)//2)
    if b == huffmanEncode.hexToBytes(target):
        if DEBUG:
            print("chromQuantTbl header good. non dynamic")
        b = f.read(64)
        chrominanceQuantTbl = np.array(list(b)).reshape([8,8])
    elif b == huffmanEncode.hexToBytes(chromQuantTbl_header_dynamic):
        if DEBUG:
            print("chromQuantTbl header good. dynamic")
        b = f.read(64)
        chrominanceQuantTbl = np.array(list(b)).reshape([8,8])
        b = f.read(64)
        chrominanceQuantTbl_lo = np.array(list(b)).reshape([8,8])
    else:
        if DEBUG:
            print("chromQuantTbl header bad")
        sys.exit()

    # height width header & value
    target = heightwidth_header
    b = f.read(len(target)//2)
    if b == huffmanEncode.hexToBytes(target):
        if DEBUG:
            print("height width header good")
        b = f.read(2)
        height = int.from_bytes(b,"big")
        b = f.read(2)
        width = int.from_bytes(b,"big")
        if DEBUG:
            print("h : {}, w : {}".format(height,width))
    else:
        if DEBUG:
            print("height width header bad")
        sys.exit()

    if (width % 8 != 0):
        width = width // 8 * 8 + 8
    if (height % 8 != 0):
        height = height // 8 * 8 + 8


    target = subsamp
    b = f.read(len(target)//2)
    if b == huffmanEncode.hexToBytes(target):
        if DEBUG:
            print("subsamp good")
    else:
        if DEBUG:
            print("subsamp bad")
        sys.exit()

    target = huffmanTable
    b = f.read(len(target)//2)
    if b == huffmanEncode.hexToBytes(target):
        if DEBUG:
            print("huffmanTable good")
    else:
        if DEBUG:
            print("huffmanTable bad")
        sys.exit()

    target = data_header
    b = f.read(len(target)//2)
    if b == huffmanEncode.hexToBytes(target):
        if DEBUG:
            print("data header good")
    else:
        if DEBUG:
            print("data header bad")
        sys.exit()

    t1 = time.time()
    print("Before loading binary : {}".format(t1 - t0))


    # load rest of binary data
    binarys = ''
    b = f.read(1)
    cnt = 1
    while 1:
        if(b == b'\xFF'):
            if f.read(1) == b'\xD9': # remove 00 after FF
                cnt = cnt - 1
                break
        bins = bin(list(b)[0])[2:]
        while len(bins) != 8:
            bins = '0' + bins
        binarys = binarys + bins
        b = f.read(1)
        cnt = cnt + 1

        if DEBUG:
            print("Binary string of full data loaded")

    t2 = time.time()
    print("Loading binary : {}".format(t2 - t1))

    # decode them into YUV matrix(zigzag ordered)
    blockSum = width // 8 * height // 8
    YZCode = np.zeros([blockSum, 64], dtype=int)
    UZCode = np.zeros([blockSum, 64], dtype=int)
    VZCode = np.zeros([blockSum, 64], dtype=int)

    # per block, dy du dv not managed yet
    yuvZZ = np.zeros([3,64],dtype=int)

    i = 1
    blockNum = 0
    zzIndex = 0
    # stage
    #   yDC_size
    #   yDC_val
    #   yAC_runsize
    #   yAC_val
    #
    #   same for u & v
    #   blockNum++ when one block of yuv is done


    stage = 0 # [Y,U,V] = [0,1,2]
    runsize = 0

    if DEBUG:
        print("blockSum: {}".format(blockSum))

    while len(binarys) != 0:
        code = list(map(int, list(binarys[:i])))
        if zzIndex == 0:
            # DC
            size = huffmanEncode.decodeDC(code, stage)
            if (size != -1):
                if DEBUG:
                    print("DC")
                    print("(stage, zzIndex, blockNum) : {}, {}, {}".format(stage, zzIndex, blockNum))
                    print("i: {}".format(i))
                    print("size: {}".format(size))
                    print("code: {}".format(code))
                binarys = binarys[i:]
                i = 1

                if(size == 0):
                    yuvZZ[stage, zzIndex] = 0
                else:
                    val = binarys[:size]
                    binarys = binarys[size:]
                    if val[0] == '1':
                        yuvZZ[stage, zzIndex] = int(val, 2)
                        if DEBUG:
                            print("val: {}".format(int(val,2)))
                    else:
                        # complement & minus sign
                        yuvZZ[stage, zzIndex] = -int(''.join(['1' if bit == '0' else '0' for bit in list(val)]),2)
                        if DEBUG:
                            print("val: {}".format(yuvZZ[stage,zzIndex]))
                zzIndex += 1
            else:
                i += 1
        else:
            # AC
            runsize = huffmanEncode.decodeAC(code, stage)
            if (runsize != None):
                if DEBUG:
                    print()
                    print("AC")
                    print("(stage, zzIndex, blockNum) : {}, {}, {}".format(stage, zzIndex, blockNum))
                    print("i: {}".format(i))
                    print("runsize: {}".format(runsize))
                    print("code: {}".format(code))

                binarys = binarys[i:]
                i = 1
                run = int(runsize[0],16)
                size = int(runsize[1],16)

                if runsize == '00':
                    while zzIndex != 64:
                        yuvZZ[stage,zzIndex] = 0
                        zzIndex += 1
                else:
                    for _ in range(run):
                        yuvZZ[stage, zzIndex] = 0
                        zzIndex += 1

                    if(size == 0):
                        yuvZZ[stage, zzIndex] = 0
                    else:
                        val = binarys[:size]
                        binarys = binarys[size:]
                        if val[0] == '1':
                            yuvZZ[stage, zzIndex] = int(val, 2)
                            if DEBUG:
                                print("val: {}".format(yuvZZ[stage,zzIndex]))
                        else:
                            # complement & minus sign
                            yuvZZ[stage, zzIndex] = -int(''.join(['1' if bit == '0' else '0' for bit in list(val)]),2)
                            if DEBUG:
                                print("val: {}".format(yuvZZ[stage,zzIndex]))
                    zzIndex += 1

                if(zzIndex == 64):
                    zzIndex = 0
                    stage += 1
                    if(stage == 3):
                        stage = 0
                        #print("blockNum: {}".format(blockNum))
                        #print("left binaries : {}".format(len(binarys)))
                        YZCode[blockNum] = yuvZZ[0]
                        UZCode[blockNum] = yuvZZ[1]
                        VZCode[blockNum] = yuvZZ[2]

                        # handle dDC
                        if(blockNum > 0):
                            YZCode[blockNum, 0] += YZCode[blockNum - 1, 0]
                            UZCode[blockNum, 0] += UZCode[blockNum - 1, 0]
                            VZCode[blockNum, 0] += VZCode[blockNum - 1, 0]

                        blockNum += 1
                        if (blockNum == blockSum):
                            if (len(binarys) <= 8):
                                if DEBUG:
                                    print("YUV of full image loaded")
                                break
                            else:
                                print("something is wrong. {} left".format(len(binarys)))
            else:
                i += 1

    # result :
    # YZCode (shape: [blockSum, 64])
    # UZCode (shape: [blockSum, 64])
    # VZCode (shape: [blockSum, 64])
    # in zig zag order.

    t3 = time.time()
    print("Main processing : {}".format(t3 - t2))

    reverseZZ = []
    for i in range(zigzagOrder.shape[0]):
        reverseZZ.append(np.where(zigzagOrder == i)[0][0])

    YCode = YZCode[:,reverseZZ]
    UCode = UZCode[:,reverseZZ]
    VCode = VZCode[:,reverseZZ]

    YQuantMatrix = YCode.reshape([blockSum,8,8])
    UQuantMatrix = UCode.reshape([blockSum,8,8])
    VQuantMatrix = VCode.reshape([blockSum,8,8])


    YDCTMatrix = np.zeros([blockSum,8,8], dtype=int)
    UDCTMatrix = np.zeros([blockSum,8,8], dtype=int)
    VDCTMatrix = np.zeros([blockSum,8,8], dtype=int)

    blockNum = 0
    for y in range(0, height, 8):
        for x in range(0, width, 8):
            if(prevDetections):
                if(isAround(bb,[x,y])):
                    YDCTMatrix[blockNum] = YQuantMatrix[blockNum] * luminanceQuantTbl
                    UDCTMatrix[blockNum] = UQuantMatrix[blockNum] * luminanceQuantTbl
                    VDCTMatrix[blockNum] = VQuantMatrix[blockNum] * luminanceQuantTbl
                else:
                    YDCTMatrix[blockNum] = YQuantMatrix[blockNum] * luminanceQuantTbl_lo
                    UDCTMatrix[blockNum] = UQuantMatrix[blockNum] * luminanceQuantTbl_lo
                    VDCTMatrix[blockNum] = VQuantMatrix[blockNum] * luminanceQuantTbl_lo
            else:
                YDCTMatrix[blockNum] = YQuantMatrix[blockNum] * luminanceQuantTbl
                UDCTMatrix[blockNum] = UQuantMatrix[blockNum] * luminanceQuantTbl
                VDCTMatrix[blockNum] = VQuantMatrix[blockNum] * luminanceQuantTbl

            blockNum += 1

    YImageMatrix = np.zeros([height, width], dtype=int)
    UImageMatrix = np.zeros([height, width], dtype=int)
    VImageMatrix = np.zeros([height, width], dtype=int)


    blockNum = 0

    for y in range(0, height, 8):
        for x in range(0, width, 8):
            YImageMatrix[y:y+8, x:x+8] = np.rint(fftpack.idct(fftpack.idct(YDCTMatrix[blockNum], axis=0, norm='ortho'), axis=1, norm='ortho'))
            UImageMatrix[y:y+8, x:x+8] = np.rint(fftpack.idct(fftpack.idct(UDCTMatrix[blockNum], axis=0, norm='ortho'), axis=1, norm='ortho'))
            VImageMatrix[y:y+8, x:x+8] = np.rint(fftpack.idct(fftpack.idct(VDCTMatrix[blockNum], axis=0, norm='ortho'), axis=1, norm='ortho'))
            blockNum += 1

    YImageMatrix = YImageMatrix + 127
    UImageMatrix = UImageMatrix + 127
    VImageMatrix = VImageMatrix + 127

    for y in range(height):
        for x in range(width):
            YImageMatrix[y,x] = max(YImageMatrix[y,x], 0)
            UImageMatrix[y,x] = max(UImageMatrix[y,x], 0)
            VImageMatrix[y,x] = max(VImageMatrix[y,x], 0)

            YImageMatrix[y,x] = min(YImageMatrix[y,x], 255)
            UImageMatrix[y,x] = min(UImageMatrix[y,x], 255)
            VImageMatrix[y,x] = min(VImageMatrix[y,x], 255)


    YImageMatrix = YImageMatrix.astype(uint8)
    UImageMatrix = UImageMatrix.astype(uint8)
    VImageMatrix = VImageMatrix.astype(uint8)

    YImage = Image.fromarray(YImageMatrix)
    UImage = Image.fromarray(UImageMatrix)
    VImage = Image.fromarray(VImageMatrix)

    decodedImage = Image.merge('YCbCr', (YImage, UImage, VImage))

    decodedImage.save(outFileName)

    t4 = time.time()
    print("After jobs : {}".format(t4 - t3))

    print("Done.")


if __name__ == '__main__':
    main()
