from PIL import Image
import math
import argparse
import os.path

#for a 2-bit dithered grayscale image use: 
# convert testimage.png -dither FloydSteinberg -define dither:diffusion-amount=85% -remap 4gray.png -compress RLE -define png:color-type=6 output.png


def swapBits(number, nbits):
    if nbits==2:
        bitstring = '{0:02b}'.format(number)
    if nbits==3:
        bitstring = '{0:03b}'.format(number)
    if nbits==4:
        bitstring = '{0:04b}'.format(number)
    if nbits==5:
        bitstring = '{0:05b}'.format(number)
    if nbits==6:
        bitstring = '{0:06b}'.format(number)
    if nbits==7:
        bitstring = '{0:07b}'.format(number)
    if nbits==8:
        bitstring = '{0:08b}'.format(number)

    bitstring=bitstring[::-1]   

    return int(bitstring, 2)


parser = argparse.ArgumentParser(
                    prog = 'Image Edit Script',
                    description = 'Manipulate or extract information from an image file',
                    epilog = '')

parser.add_argument('filename')           # positional argument
parser.add_argument('-n', '--bytesperline')
parser.add_argument('-d', '--bitdepth')
parser.add_argument('-b', '--lsbfirst', action='store_true') #lsb first for the whole byte
parser.add_argument('-c', '--colorlsbfirst', action='store_true') #lsb first for a single color
parser.add_argument('-i', '--invert', action='store_true')

args = parser.parse_args()

lsbfirst=args.lsbfirst
colorlsbfirst=args.colorlsbfirst
invert=args.invert


bitdepth=1
if args.bitdepth is not None:
    _bd=int(int(args.bitdepth))
    if _bd!=1 and _bd!=2 and _bd!=4 and _bd!=8:
        print("bitdepth needs to be 1, 2, 4, or 8")
        exit()
    bitdepth=_bd

bytesperline=16
if args.bytesperline is not None:
    bytesperline=int(args.bytesperline) #for output formatting


im = Image.open(args.filename) # Can be many different formats.
pix = im.load()
print(im.size)  # Get the width and hight of the image for iterating over


array=[] #array with every element a byte

    
def calculateDistance(x1,y1,x2,y2):  
    dist = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)  
    return dist  
    
temp_byte=0
temp_byte_pos=0
for y in range(im.size[1]):
    for x in range(im.size[0]):
        c = pix[x,y] #get pixel
        r=c[0]
        g=c[1]
        b=c[2]
        
        _a=int((r+g+b)/3 )
        brightness=int(_a)

        brightnessreduced=brightness>>(8-bitdepth) #reduce to bitdepth

        brightnessreduced_byte=brightnessreduced<<(8-bitdepth) #scale back to byte
        pix[x,y]=(brightnessreduced_byte,brightnessreduced_byte,brightnessreduced_byte)

        
        if invert:
            brightnessreduced=pow(2,bitdepth)-1-brightnessreduced #only invert txt output

        if colorlsbfirst:
            brightnessreduced=swapBits(brightnessreduced,bitdepth)


        temp_byte+=brightnessreduced<<temp_byte_pos #add to current byte
        
        

        temp_byte_pos+=bitdepth
        #print("temp byte "+str(temp_byte)+" pos="+str(temp_byte_pos)+" brightnessreduced="+str(brightnessreduced))
        if temp_byte_pos>=8: #finished assemblying byte
            if lsbfirst: #swap bit order
                temp_byte=swapBits(temp_byte,8)
            array.append(temp_byte)
            #print("appending "+str(temp_byte))
            temp_byte_pos=0 #reset
            temp_byte=0


im.save(args.filename+'_result'+'.png')

        
if os.path.isfile(args.filename+'.txt'):
    print("Outputfile "+args.filename+".txt exists")
    exit()

with open(args.filename+'.txt', 'w') as f:
    f.write("const unsigned char gImage_4in2[]  = {")
    f.write('\r\n')
    counter=0
    for a in array: #for every byte
        hexstring="0X{:02X}".format(a)
        f.write(hexstring+',') #Example output: 0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,0XFF,
        counter+=1
        if counter>=bytesperline:
            f.write('\r\n')
            counter=0
    f.write("};")