import networkx as nx
from PIL import Image, ImageDraw, ImageFont
import mahotas as mh

def branchedPoints(skel, showSE=True):
    X=[]
    #cross X
    X0 = np.array([[0, 1, 0], 
                   [1, 1, 1], 
                   [0, 1, 0]])
    X1 = np.array([[1, 0, 1], 
                   [0, 1, 0], 
                   [1, 0, 1]])
    X.append(X0)
    X.append(X1)
    #T like
    T=[]
    #T0 contains X0
    T0=np.array([[2, 1, 2], 
                 [1, 1, 1], 
                 [2, 2, 2]])
            
    T1=np.array([[1, 2, 1], 
                 [2, 1, 2],
                 [1, 2, 2]])  # contains X1
  
    T2=np.array([[2, 1, 2], 
                 [1, 1, 2],
                 [2, 1, 2]])
    
    T3=np.array([[1, 2, 2],
                 [2, 1, 2],
                 [1, 2, 1]])
    
    T4=np.array([[2, 2, 2],
                 [1, 1, 1],
                 [2, 1, 2]])
    
    T5=np.array([[2, 2, 1], 
                 [2, 1, 2],
                 [1, 2, 1]])
    
    T6=np.array([[2, 1, 2],
                 [2, 1, 1],
                 [2, 1, 2]])
    
    T7=np.array([[1, 2, 1],
                 [2, 1, 2],
                 [2, 2, 1]])
    T.append(T0)
    T.append(T1)
    T.append(T2)
    T.append(T3)
    T.append(T4)
    T.append(T5)
    T.append(T6)
    T.append(T7)
    #Y like
    Y=[]
    Y0=np.array([[1, 0, 1], 
                 [0, 1, 0], 
                 [2, 1, 2]])
    
    Y1=np.array([[0, 1, 0], 
                 [1, 1, 2], 
                 [0, 2, 1]])
    
    Y2=np.array([[1, 0, 2], 
                 [0, 1, 1], 
                 [1, 0, 2]])
    
    Y2=np.array([[1, 0, 2], 
                 [0, 1, 1], 
                 [1, 0, 2]])
    
    Y3=np.array([[0, 2, 1], 
                 [1, 1, 2], 
                 [0, 1, 0]])
    
    Y4=np.array([[2, 1, 2], 
                 [0, 1, 0], 
                 [1, 0, 1]])
    Y5=np.rot90(Y3)
    Y6 = np.rot90(Y4)
    Y7 = np.rot90(Y5)
    Y.append(Y0)
    Y.append(Y1)
    Y.append(Y2)
    Y.append(Y3)
    Y.append(Y4)
    Y.append(Y5)
    Y.append(Y6)
    Y.append(Y7)
    
    bp = np.zeros(skel.shape, dtype=int)
    for x in X:
        bp = bp + mh.morph.hitmiss(skel,x)
    for y in Y:
        bp = bp + mh.morph.hitmiss(skel,y)
    for t in T:
        bp = bp + mh.morph.hitmiss(skel,t)
        
    if showSE==True:
        fig = plt.figure(figsize=(4,5))
        tX =['X0','X1']
        tY =['Y'+str(i) for i in range(0,8)]
        tT =['T'+str(i) for i in range(0,8)]
        ti= tX+tY+tT
        SE=X+Y+T
        print len(SE), len(ti)
        n = 1
        ti = iter(ti)
        for se in SE:
            #print next(ti)
            #print se
            mycmap = mpl.colors.ListedColormap(['black','blue','red'])
            ax = fig.add_subplot(4,5,n,frameon=False, xticks=[], yticks=[])
            title(str(next(ti)))
            imshow(se, interpolation='nearest',vmin=0,vmax=2,cmap=mycmap)
            n = n+1
        fig.subplots_adjust(hspace=0.1,wspace=0.08)
        #ax_cb = fig.add_axes([.9,.25,.1,.3])#
        color_vals=[0,1,2]
        #cb = mpl.colorbar.ColorbarBase(ax_cb,cmap=mycmap, ticks=color_vals)
        #cb.set_ticklabels(['back', 'hit', 'don\'t care'])
        
        plt.show()
    return bp > 0

def endPoints(skel):
    endpoint1=np.array([[0, 0, 0],
                        [0, 1, 0],
                        [2, 1, 2]])
    
    endpoint2=np.array([[0, 0, 0],
                        [0, 1, 2],
                        [0, 2, 1]])
    
    endpoint3=np.array([[0, 0, 2],
                        [0, 1, 1],
                        [0, 0, 2]])
    
    endpoint4=np.array([[0, 2, 1],
                        [0, 1, 2],
                        [0, 0, 0]])
    
    endpoint5=np.array([[2, 1, 2],
                        [0, 1, 0],
                        [0, 0, 0]])
    
    endpoint6=np.array([[1, 2, 0],
                        [2, 1, 0],
                        [0, 0, 0]])
    
    endpoint7=np.array([[2, 0, 0],
                        [1, 1, 0],
                        [2, 0, 0]])
    
    endpoint8=np.array([[0, 0, 0],
                        [2, 1, 0],
                        [1, 2, 0]])
    
    ep1=mh.morph.hitmiss(skel,endpoint1)
    ep2=mh.morph.hitmiss(skel,endpoint2)
    ep3=mh.morph.hitmiss(skel,endpoint3)
    ep4=mh.morph.hitmiss(skel,endpoint4)
    ep5=mh.morph.hitmiss(skel,endpoint5)
    ep6=mh.morph.hitmiss(skel,endpoint6)
    ep7=mh.morph.hitmiss(skel,endpoint7)
    ep8=mh.morph.hitmiss(skel,endpoint8)
    ep = ep1+ep2+ep3+ep4+ep5+ep6+ep7+ep8
    return ep > 0

def pruning(skeleton, size):
    '''remove iteratively end points "size" 
       times from the skeleton
    '''
    for i in range(0, size):
        endpoints = endPoints(skeleton)
        endpoints = np.logical_not(endpoints)
        skeleton = np.logical_and(skeleton,endpoints)
    return skeleton

def edges_from_C8skel(c8skeleton):
    '''given a skeleton defined on c8 neighborhood (use mahotas),
     \ returns labeled edges
    '''
    branchedP = branchedPoints(c8skeleton, showSE = False) > 0
    endP = endPoints(c8skeleton) > 0
    edges = np.logical_not(branchedP)*c8skeleton
    label_edges,ne = mh.label(edges)
    return label_edges

def add_bp_to_graph(Graph, labels_branchpoints):
    labels = np.arange(1, labels_branchpoints.max()+1)
    label_bp = {}
    for lab in labels:
        pos = np.where(labels_branchpoints==lab)
        row = int(pos[0])
        col = int(pos[1])
        label_bp[lab] = lab
        Graph.add_node(lab, pos= (row,col), label=lab)
    return label_bp

def add_ep_to_graph(Graph, im_labels_endpoints):
    label_ep = {}
    first_endpoint = Graph.number_of_nodes()+1
    labels = np.where(mh.histogram.fullhistogram(np.uint16(im_labels_endpoints))[:]==1)[0]
    n=0
    for lab in labels:
        pos = np.where(im_labels_endpoints==lab)
        row = int(pos[0])
        col = int(pos[1])
        nodeIndex = first_endpoint + n
        Graph.add_node( nodeIndex, pos= (row,col), label=lab)
        label_ep[lab]= nodeIndex
        n = n+1
    return label_ep

def C8skeleton_to_graph(skeletonC8):
    #images processing: extraction of branchedpoints, end-points, edges
    ep = endPoints(skeletonC8)
    bp = branchedPoints(skeletonC8, showSE = False)
    ## Label branched-points
    l_bp,_ = mh.label(bp)
    ## Label the edges
    l_edges =  edges_from_C8skel(skeletonC8)
    ##Make end-points with the same label than their edge
    l_ep = ep*l_edges 
    
    ##edges between branched-points
    endpoints_labels = np.where(mh.histogram.fullhistogram(np.uint16(l_ep))[:]==1)[0]
    edges_bp = np.copy(l_edges)
    for l in endpoints_labels:
        edges_bp[np.where(edges_bp == l)]=0
    #edges between end-points and branched points
    edges_ep_to_bp = l_edges * np.logical_not(edges_bp > 0)
    
    #Building graph
    ## Add branched points first
    G=nx.Graph()
    lab_bpTonode = add_bp_to_graph(G, l_bp)
    ## Add end-points
    lab_epTonode = add_ep_to_graph(G, l_ep)
    ##Link end-points to branched-points
    ###labels of bp
    branchedpoints_labels = np.where(mh.histogram.fullhistogram(np.uint16(l_bp))[:]==1)[0]
    for lab in branchedpoints_labels:
        pos = np.where(l_bp == lab)
        row = int(pos[0])
        col = int(pos[1])
        #search label(s) of edges in image containing edges between ep and bp
        ## first get the neighborhood of the curent bp
        neigh_epbp = edges_ep_to_bp[row-1:row+1+1,col-1:col+1+1]
        labels_in_neigh = np.where(mh.histogram.fullhistogram(np.uint16(neigh_epbp))[:]<>0)[0]
        #print neigh_epbp, labels_in_neigh[1:]
        #get node(s) of attribute label= labels_in_neigh ! may be more than one, think to a  list
        for lab_ep in labels_in_neigh[1:]:
            #search for nodes f attribute label= lab_ep
            w = np.sum(l_edges==lab_ep)
            #print 'linking ',lab, lab_ep, ' weight ',w
            G.add_edge(lab_bpTonode[lab],lab_epTonode[lab_ep],weight=w)#     
    ##
    ##Now try to link branched points between them
    ##
    bps_neighborhood = {}
    branchedpoints_labels = np.where(mh.histogram.fullhistogram(np.uint16(l_bp))[:]==1)[0]
    for lab in branchedpoints_labels:
        pos = np.where(l_bp == lab)
        row = int(pos[0])
        col = int(pos[1])
        #search label(s) of edges in image containing edges between ep and bp
        ## first get the neighborhood of the curent bp
        neigh_epbp = edges_bp[row-1:row+1+1,col-1:col+1+1]
        labels_in_neigh = np.where(mh.histogram.fullhistogram(np.uint16(neigh_epbp))[:]<>0)[0]
        bps_neighborhood[lab]=labels_in_neigh[1:].tolist()
    #print bps_neighborhood
    
    ## Build the dictionnary of edges see (http://stackoverflow.com/questions/21375146/pythonic-inverse-dict-non-unique-mappings)
    invert_is_edges = {item: [key for key in bps_neighborhood if item in bps_neighborhood[key]] for value in bps_neighborhood.values() for item in value}
    ## Addeges to graph
    for ed in invert_is_edges.keys():
        ## first get edge size -> its weight
        w = np.sum(edges_bp==ed)
        vertex1 = invert_is_edges[ed][0]
        vertex2 = invert_is_edges[ed][1]
        #print ed,w
        G.add_edge(vertex1,vertex2,weight=w)
    ## This is it !!
    return G

def makeLetterImage(character, size):
    image = Image.new("RGBA", (600,150), (255,255,255))
    draw = ImageDraw.Draw(image)
    font = ImageFont.truetype("verdana.ttf", size)
    
    draw.text((5, 0), character, (0,0,0), font=font)
    img_resized = np.array(image.resize((300,75), Image.ANTIALIAS))
    letter = img_resized[0:40,0:30,0]<64
    r1 = mh.bbox(letter)[0]
    r2 = mh.bbox(letter)[1]
    c1 = mh.bbox(letter)[2]
    c2 = mh.bbox(letter)[3]
    letter = letter[r1-1:r2+1,c1-1:c2+1]
    return letter

imA = makeLetterImage("A", 70)
imB = makeLetterImage("B", 70)
imH = makeLetterImage("H", 70)
imO= makeLetterImage("O", 70)

skeletonB = mh.thin(imB)
Bgraph = C8skeleton_to_graph(skeletonB)

skeletonA = mh.thin(imA)
Agraph = C8skeleton_to_graph(skeletonA)

skeletonO = mh.thin(imO)
Og = C8skeleton_to_graph(skeletonO)

figsize(6,6)
subplot(221,xticks=[],yticks=[])
imshow(imA, interpolation='nearest')
subplot(222,xticks=[],yticks=[])
nx.draw(Agraph)
subplot(223,xticks=[],yticks=[])
imshow(skeletonO, interpolation='nearest')
subplot(224,xticks=[],yticks=[])
nx.draw(Og)

import string

alphabet = list(string.ascii_uppercase)

ilphabet = []
for l in alphabet:
    imLetter = makeLetterImage(l, 70)
    skeletter = mh.thin(imLetter)
    BP =  branchedPoints(skeletter, showSE=False)
    EP = endPoints(skeletter)
    ilphabet.append(skeletter+10*BP+3*EP)#np.uint(skeletter) +

figsize(10,10)
for im,n in zip(ilphabet, range(1,27)):
    subplot(6,5,n,xticks=[],yticks=[])
    imshow(im,interpolation='nearest')

gralphabet = []
ilphabet = []
for l in alphabet:
    imLetter = makeLetterImage(l, 70)
    skeletter = mh.thin(imLetter)
    print l
    gralphabet.append(C8skeleton_to_graph(skeletter))
    

print len(gralphabet)
figsize(10,15)
for g,n in zip(gralphabet, range(1,11)):
    subplot(4,3,n)
    nx.draw(g)