from matplotlib import pyplot as plt
from matplotlib import image as img
import numpy as np
import glob, os

shapes  = ['D', 'T', 'C', 'S']
threshs = [20, 125, 180, 235]
bandwidth = 18

def determine_symbol(inCnt):
    x = max([thresh for thresh in threshs if inCnt > thresh])
    return shapes[threshs.index(x)]

def strornone(x):
    if x:
        return ''.join(x)
    else:
        return None

# pass in the image filename when you declare the class instance
class puzzleImage():
    def __init__(self, imgFile, verbose=False):
        self.image = img.imread(imgFile).astype('int')
        self.parseSymbols(verbose)
        self.filename = imgFile
        self.setcoord(0,0)

    def setcoord(self,x,y):
        self.x=x
        self.y=y

    def parseSymbols(self,verbose=False):
        #start on left side column, track which row you see non-white pixels on
        #and which row you go back to white pixels on, this will allow us to box each symbol.
        myimage = self.image

        boxtops = []
        boxbots = []

        switch = 0
        for row in range(myimage.shape[0]):
            whitecnt = np.count_nonzero( np.amax(myimage[row,:bandwidth], axis=1) > 0xF0)
            #print(whitecnt)
            if (switch == 0) and (whitecnt < bandwidth - 4):
                boxtops.append(row)
                switch = 1
            if (switch == 1) and (whitecnt > bandwidth - 2):
                boxbots.append(row-1)
                switch = 0
                if abs(boxbots[-1] - boxtops[-1]) < 5:
                    boxbots = boxbots[:-1]
                    boxtops = boxtops[:-1]
            if len(boxbots) == 5:
                break

        if len(boxbots) < 5:
            L_symbols = None
        else:
            black_pixel_cnts = [np.count_nonzero( np.amax(myimage[boxtops[j]:boxbots[j],:bandwidth], axis=2) < 0xF0) for j in range(5)]
            L_symbols = [determine_symbol(x) for x in black_pixel_cnts]
            if verbose: print(black_pixel_cnts)
        #Now right side

        boxtops = []
        boxbots = []

        switch = 0
        for row in range(myimage.shape[0]):
            whitecnt = np.count_nonzero( np.amax(myimage[row,-bandwidth:], axis=1) > 0xF0)

            if (switch == 0) and (whitecnt < bandwidth - 4):
                boxtops.append(row)
                switch = 1
            if (switch == 1) and (whitecnt > bandwidth - 2):
                boxbots.append(row-1)
                switch = 0
                if abs(boxbots[-1] - boxtops[-1]) < 5:
                    boxbots = boxbots[:-1]
                    boxtops = boxtops[:-1]
            if len(boxbots) == 5:
                break

        if len(boxbots) < 5:
            R_symbols = None
        else:
            black_pixel_cnts = [np.count_nonzero( np.amax(myimage[boxtops[j]:boxbots[j],-bandwidth:], axis=2) < 0xF0) for j in range(5)]
            R_symbols = [determine_symbol(x) for x in black_pixel_cnts]
            if verbose: print(black_pixel_cnts)

        #Now top

        boxtops = []
        boxbots = []

        switch = 0
        for col in range(myimage.shape[1]):
            whitecnt = np.count_nonzero( np.amax(myimage[:bandwidth,col], axis=1) > 0xF0)

            if (switch == 0) and (whitecnt < bandwidth - 4):
                boxtops.append(col)
                switch = 1
            if (switch == 1) and (whitecnt > bandwidth - 2):
                boxbots.append(col-1)
                switch = 0
                if abs(boxbots[-1] - boxtops[-1]) < 5:
                    boxbots = boxbots[:-1]
                    boxtops = boxtops[:-1]
            if len(boxbots) == 5:
                break

        if len(boxbots) < 5:
            T_symbols = None
        else:
            black_pixel_cnts = [np.count_nonzero( np.amax(myimage[:bandwidth, boxtops[j]:boxbots[j]], axis=2) < 0xF0) for j in range(5)]
            T_symbols = [determine_symbol(x) for x in black_pixel_cnts]
            if verbose: print(black_pixel_cnts)

        #Now bot

        boxtops = []
        boxbots = []

        switch = 0
        for col in range(myimage.shape[1]):
            whitecnt = np.count_nonzero( np.amax(myimage[-bandwidth:,col], axis=1) > 0xF0)

            if (switch == 0) and (whitecnt < bandwidth - 4):
                boxtops.append(col)
                switch = 1
            if (switch == 1) and (whitecnt > bandwidth - 2):
                boxbots.append(col-1)
                switch = 0
                if abs(boxbots[-1] - boxtops[-1]) < 5:
                    boxbots = boxbots[:-1]
                    boxtops = boxtops[:-1]
            if len(boxbots) == 5:
                break

        if len(boxbots) < 5:
            B_symbols = None
        else:
            black_pixel_cnts = [np.count_nonzero( np.amax(myimage[-bandwidth:, boxtops[j]:boxbots[j]], axis=2) < 0xF0) for j in range(5)]
            B_symbols = [determine_symbol(x) for x in black_pixel_cnts]
            if verbose: print(black_pixel_cnts)


        self.Lsym = strornone(L_symbols)
        self.Rsym = strornone(R_symbols)
        self.Bsym = strornone(B_symbols)
        self.Tsym = strornone(T_symbols)

        #print(self.Lsym, self.Rsym, self.Bsym, self.Tsym)

# actually doesn't print the chunk out any more
# instead pastes together a set of images using their coords.
def printchunk(expset):
    minx = min([elt.x for elt in expset])
    maxx = max([elt.x for elt in expset])
    miny = min([elt.y for elt in expset])
    maxy = max([elt.y for elt in expset])

    for elt in expset:
        elt.setcoord(elt.x - minx, elt.y - miny)

    maxx -= minx
    maxy -= miny

    big = np.ones((130*(maxx+1),130*(maxy+1),3))*255
    print(big.shape, elt.image.shape)
    for elt in expset:
        big[130*elt.x : 130*elt.x + 130, 130*elt.y:130*elt.y + 130] = elt.image[25:-25, 25:-25]

    return big

if __name__ == '__main__':
    imagefiles = glob.glob('*JPEG')
    print('Found {0} image files'.format(len(imagefiles)))
    parsedimages = [puzzleImage(x) for x in imagefiles]

    expandingsets = [[x] for x in parsedimages]

    prev = len(expandingsets)
    new = 0

    while new < prev:
        prev = len(expandingsets)
        for i in range(len(expandingsets)):
            remlist = []
            expset1 = expandingsets[i]
            for j in range(i+1, len(expandingsets)):
                expset2 = expandingsets[j]

                #TSYM
                if set([x.Tsym for x in expset1 if x.Tsym]) & set([x.Rsym for x in expset2 if x.Rsym]):
                    remlist.append(j)
                    cstr = list(set([x.Tsym for x in expset1 if x.Tsym]) & set([x.Rsym for x in expset2 if x.Rsym]))[0]
                    chosen1 = [q for q in expset1 if q.Tsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Rsym == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x - 1 - chosen2.x, q.y + chosen1.y - 1 - chosen2.y)
                    expset1 += expset2
                elif set([x.Tsym for x in expset1 if x.Tsym]) & set([x.Bsym[::-1] for x in expset2 if x.Bsym]):
                    remlist.append(j)
                    cstr = list(set([x.Tsym for x in expset1 if x.Tsym]) & set([x.Bsym[::-1] for x in expset2 if x.Bsym]))[0]
                    chosen1 = [q for q in expset1 if q.Tsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Bsym[::-1] == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x - 2 - chosen2.x, q.y + chosen1.y - chosen2.y)
                    expset1 += expset2
                elif set([x.Tsym for x in expset1 if x.Tsym]) & set([x.Lsym[::-1] for x in expset2 if x.Lsym]):
                    remlist.append(j)
                    cstr = list(set([x.Tsym for x in expset1 if x.Tsym]) & set([x.Lsym[::-1] for x in expset2 if x.Lsym]))[0]
                    chosen1 = [q for q in expset1 if q.Tsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Lsym[::-1] == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x - 1 - chosen2.x, q.y + chosen1.y + 1 - chosen2.y)
                    expset1 += expset2

                #RSYM
                elif set([x.Rsym for x in expset1 if x.Rsym]) & set([x.Tsym for x in expset2 if x.Tsym]):
                    remlist.append(j)
                    cstr = list(set([x.Rsym for x in expset1 if x.Rsym]) & set([x.Tsym for x in expset2 if x.Tsym]))[0]
                    chosen1 = [q for q in expset1 if q.Rsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Tsym == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x + 1 - chosen2.x, q.y + chosen1.y + 1 - chosen2.y)
                    expset1 += expset2
                elif set([x.Rsym for x in expset1 if x.Rsym]) & set([x.Lsym[::-1] for x in expset2 if x.Lsym]):
                    remlist.append(j)
                    cstr = list(set([x.Rsym for x in expset1 if x.Rsym]) & set([x.Lsym[::-1] for x in expset2 if x.Lsym]))[0]
                    chosen1 = [q for q in expset1 if q.Rsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Lsym[::-1] == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x - chosen2.x, q.y + chosen1.y + 2 - chosen2.y)
                    expset1 += expset2
                elif set([x.Rsym for x in expset1 if x.Rsym]) & set([x.Bsym[::-1] for x in expset2 if x.Bsym]):
                    remlist.append(j)
                    cstr = list(set([x.Rsym for x in expset1 if x.Rsym]) & set([x.Bsym[::-1] for x in expset2 if x.Bsym]))[0]
                    chosen1 = [q for q in expset1 if q.Rsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Bsym[::-1] == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x - 1 - chosen2.x, q.y + chosen1.y + 1 - chosen2.y)
                    expset1 += expset2


                #LSYM
                elif set([x.Lsym for x in expset1 if x.Lsym]) & set([x.Tsym[::-1] for x in expset2 if x.Tsym]):
                    remlist.append(j)
                    cstr = list(set([x.Lsym for x in expset1 if x.Lsym]) & set([x.Tsym[::-1] for x in expset2 if x.Tsym]))[0]
                    chosen1 = [q for q in expset1 if q.Lsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Tsym[::-1] == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x + 1 - chosen2.x, q.y + chosen1.y - 1 - chosen2.y)
                    expset1 += expset2
                elif set([x.Lsym for x in expset1 if x.Lsym]) & set([x.Rsym[::-1] for x in expset2 if x.Rsym]):
                    remlist.append(j)
                    cstr = list(set([x.Lsym for x in expset1 if x.Lsym]) & set([x.Rsym[::-1] for x in expset2 if x.Rsym]))[0]
                    chosen1 = [q for q in expset1 if q.Lsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Rsym[::-1] == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x - chosen2.x, q.y + chosen1.y - 2 - chosen2.y)
                    expset1 += expset2
                elif set([x.Lsym for x in expset1 if x.Lsym]) & set([x.Bsym for x in expset2 if x.Bsym]):
                    remlist.append(j)
                    cstr = list(set([x.Lsym for x in expset1 if x.Lsym]) & set([x.Bsym for x in expset2 if x.Bsym]))[0]
                    chosen1 = [q for q in expset1 if q.Lsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Bsym == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x - 1 - chosen2.x, q.y + chosen1.y - 1 - chosen2.y)
                    expset1 += expset2


                #BSYM
                elif set([x.Bsym for x in expset1 if x.Bsym]) & set([x.Rsym[::-1] for x in expset2 if x.Rsym]):
                    remlist.append(j)
                    cstr = list(set([x.Bsym for x in expset1 if x.Bsym]) & set([x.Rsym[::-1] for x in expset2 if x.Rsym]))[0]
                    chosen1 = [q for q in expset1 if q.Bsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Rsym[::-1]== cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x + 1 - chosen2.x, q.y + chosen1.y + 1 - chosen2.y)
                    expset1 += expset2
                elif set([x.Bsym for x in expset1 if x.Bsym]) & set([x.Tsym[::-1] for x in expset2 if x.Tsym]):
                    remlist.append(j)
                    cstr = list(set([x.Bsym for x in expset1 if x.Bsym]) & set([x.Tsym[::-1] for x in expset2 if x.Tsym]))[0]
                    chosen1 = [q for q in expset1 if q.Bsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Tsym[::-1]== cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x + 2 - chosen2.x, q.y + chosen1.y - chosen2.y)
                    expset1 += expset2
                elif set([x.Bsym for x in expset1 if x.Bsym]) & set([x.Lsym for x in expset2 if x.Lsym]):
                    remlist.append(j)
                    cstr = list(set([x.Bsym for x in expset1 if x.Bsym]) & set([x.Lsym for x in expset2 if x.Lsym]))[0]
                    chosen1 = [q for q in expset1 if q.Bsym == cstr][0]
                    chosen2 = [q for q in expset2 if q.Lsym == cstr][0]
                    for q in expset2:
                        q.setcoord(q.x + chosen1.x + 1 - chosen2.x, q.y + chosen1.y + 1 - chosen2.y)
                    expset1 += expset2

            remlist.reverse()
            for j in remlist:
                expandingsets.pop(j)
            if remlist:
                break

        new = len(expandingsets)
        print(new)

    bigs = []
    for elt in expandingsets:
        if len(elt) > 1:
            bigs.append(printchunk(elt))

	#I knew there were only 2, could have done this more generally
    bigall = bigs[0] + bigs[1]

    im1 = bigs[0]/255.
    im2 = bigs[1]/255.
    im3 = bigall/255.

    img.imsave('testplot1.jpg',im1,format='JPEG')
    img.imsave('testplot2.jpg',im2,format='JPEG')
    img.imsave('testplot3.jpg',im3,format='JPEG')
    plt.imshow(bigall/255.)
    plt.show()