import numpy as np
import matplotlib.pyplot as plt
from skimage.segmentation import mark_boundaries
import os
import pickle as pkl
import numpy.random as rand
import matplotlib.lines as lines


def drawGraph(graphe, im, seg, coord):
    fig,ax=plt.subplots()
    color = ['blue','red','green','black']
    ax.imshow(mark_boundaries(im, seg),interpolation='nearest')
    scale = len(im)
    ax.scatter(coord[:,2]*scale, coord[:,0]*scale)
    for i,node in enumerate(graphe):
        for c,side in enumerate(node):
            if c==1 or c ==3:
                continue
            #print side
            if not side == -1:
                ax.add_line(lines.Line2D(coord[[i,side],2]*scale, coord[[i,side],0]*scale, linewidth=1, color=color[c]))
    #ax.set_title('graphe de l\'image')
    return

def drawTree(graphe, tree, im, seg, coord, bg=True):
    fig, ax = plt.subplots()
    if bg:
        ax.imshow(mark_boundaries(im, seg),interpolation='nearest')
    else:
        ax.imshow(mark_boundaries(im, seg),interpolation='nearest', alpha=0.8)
    ax.scatter(coord[:,2]*len(im), coord[:,0]*len(im))
    for i,node in enumerate(graphe):
        #for c,side in enumerate(node):
        if tree[i] == 1:
            if not node[2] == -1:
                ax.add_line(lines.Line2D(coord[[i,node[2]],2]*len(im), coord[[i,node[2]],0]*len(im), linewidth=1, color=color[0]))
        else:
            if not node[0] == -1:
                ax.add_line(lines.Line2D(coord[[i,node[0]],2]*len(im), coord[[i,node[0]],0]*len(im), linewidth=1, color=color[1]))
    #lines.Line2D(rand.rand(2), rand.rand(2))
    return