def make_venn(direc,file1,file2,filename,text1,text2,color1,color2,color3):
    smallername = ""
    biggername = ""
    fileinput1 = direc + file1
    fileinput2 = direc + file2
    statinfo1 = os.stat(fileinput1)
    statinfo2 = os.stat(fileinput2)
    if statinfo1.st_size > statinfo2.st_size:
        smallerfile = open(fileinput2)
        biggerfile = open(fileinput1)
        smallername = text2
        biggername = text1
    else:
        smallerfile = open(fileinput1)
        biggerfile = open(fileinput2)
        smallername = text1
        biggername = text2
    overlap_info = peak_overlap(smallerfile, biggerfile)
    file1_nums = overlap_info[1] - overlap_info[0]
    file2_nums = overlap_info[2] - overlap_info[0]
    print file1_nums, file2_nums
    if biggername == text1:
        v= venn2(subsets = (file2_nums, file1_nums, overlap_info[0]), set_labels = (biggername, smallername))
    else:
        v= venn2(subsets = (file1_nums, file2_nums, overlap_info[0]), set_labels = (smallername, biggername))
    print file2_nums, file1_nums, overlap_info[0]
    v.get_patch_by_id('10').set_color(color1)
    v.get_patch_by_id('01').set_color(color2)
    v.get_patch_by_id('11').set_color(color3)
    v.get_patch_by_id('10').set_alpha(1.0)
    v.get_patch_by_id('01').set_alpha(1.0)
    v.get_patch_by_id('11').set_alpha(0.7)
    plt.title(filename)
    savefig(filename)
    plt.close()

def plot_venn(t1=None, t2=None, t3=None, ax=None, set_colors=('r', 'b', 'k')):
	"""input: 2 or 3 tuples: (list/set, name_to_display) """
	assert len(t1) == len(t2) == 2
	if t3:
		venn3( [set(t[0]) for t in [t1,t2,t3]], tuple( ['%s\n(%s)'%(t[1], len(set(t[0])) ) for t in [t1,t2,t3]]) , set_colors=set_colors, alpha=0.5,ax=ax)
	else:
		venn2( [set(t[0]) for t in [t1,t2]], tuple( ['%s\n(%s)'%(t[1], len(set(t[0])) ) for t in [t1,t2]]), set_colors=set_colors[0:2],alpha=0.5, ax=ax)

def compre_two(ans_file1, ans_file2):

    with open(ans_file1) as f:
        a = json.load(f)
    with open(ans_file2) as f:
        b = json.load(f)

    true_a = []
    wrong_a = []
    for ele in a:
        if ele['em']:
            true_a.append(ele['id'])
        if not ele['f1']:
            wrong_a.append(ele['id'])

    true_b = []
    wrong_b = []
    for ele in b:
        if not ele['f1']:
            wrong_b.append(ele['id'])
        if ele['em']:
            true_b.append(ele['id'])

    label_a = get_label(ans_file1)
    label_b = get_label(ans_file2)
    venn2([set(true_a), set(true_b)], set_labels=(label_a, label_b))

    diff = set(true_b).intersection(set(wrong_a))
    if len(diff) > 20:
        diff = list(diff)[:20]
    print('true in b, but wrong in a:')
    print(diff)

def venn_pair( set1, set2, mask):
    set_a = set([v.strip() for v in open( mask % ethnicity_code[set1]).readlines()])
    set_b = set([v.strip() for v in open( mask % ethnicity_code[set2]).readlines()])
    
    plt.cla()
    venn2((set_a, set_b), set_labels=(set1, set2))
    plt.savefig("%s_%s_%s.png" % (set1, set2, 'venn'+mask.split('%s')[1] ))

def draw_venn_diag(venn_dataset, first_name, second_name, out_dir):
    """Draws Venn diagram comparing genes hit in each dataset.

    Parameters
    ----------
    venn_dataset    :   list of integers
        First element   :   Number of genes with hits in first dataset only
        Second element  :   Number of genes with hits in second dataset only
        Third element   :   Number of genes with hits in both datasets
    *_name  :   string
        Filename for first or second input hits file directory
    out_dir :   string
        Name for output directory

    Writes
    -------
    venn_diag.*.png :   png file
        Venn diagram image
    """
    venn2(subsets=(venn_dataset[0], venn_dataset[1], venn_dataset[2]), 
            set_labels = (first_name, second_name))
    
    plt.title('$C.$ $albicans$ genes with hits')
    plt.savefig(os.path.join(out_dir, 'venn_diag.%s_and_%s.png' % (first_name, second_name)))
    plt.close()

    def func_plot( self, json_data, str_output_figure ):
        """
        Function that quickly plots a venn diagram of data in a json file.
        """
        
        str_title = json_data[ qp.c_STR_TITLE ] if qp.c_STR_TITLE in json_data else qp.c_STR_TITLE_DEFAULT
        ldict_data = json_data[ qp.c_STR_DATA ]

        # Two venn diagram mode
        if len( ldict_data ) == 2:
            # Plot venn diagram
            # Labels
            str_data_label_1 = ldict_data[ 0 ][ qp.c_STR_DATA_LABEL ] if qp.c_STR_DATA_LABEL in ldict_data[ 0 ] else None
            str_data_label_2 = ldict_data[ 1 ][ qp.c_STR_DATA_LABEL ] if qp.c_STR_DATA_LABEL in ldict_data[ 1 ] else None
            # Set colors
            str_data_color_1 = ldict_data[ 0 ][ qp.c_C_PLOT_COLOR ] if qp.c_C_PLOT_COLOR in ldict_data[ 0 ] else 'r'
            str_data_color_2 = ldict_data[ 1 ][ qp.c_C_PLOT_COLOR ] if qp.c_C_PLOT_COLOR in ldict_data[ 1 ] else 'g'
            venn2( [ set( ldict_data[ 0 ][ qp.c_STR_DATA ] ), set( ldict_data[ 1 ][ qp.c_STR_DATA ] ) ],
                set_labels = [ str_data_label_1, str_data_label_2 ], set_colors = [ str_data_color_1, str_data_color_2 ] )
        else:
            return False

        plt.title( str_title )
        plt.tight_layout()
        plt.savefig( str_output_figure )
        plt.close()
        
        return True

def venn_diagram(df1, df2, labels=['A', 'B'], save=False):
    from matplotlib_venn import venn2
    import pandas as pd

    try:
        labels = [s.strip('.txt') for s in labels]
    except:
        pass

    s1 = set(df1.index)
    s2 = set(df2.index)
    venn2([s1, s2], set_colors=['navy','lightblue'], set_labels=labels)

    if save:
        inter = [s for s in s1 & s2]
        s1_only = [s for s in s1 - s2]
        s2_only = [s for s in s2 - s1]
        merge = pd.merge(df1, df2, how='outer',
                         left_index=True, right_index=True)
        merge.loc[inter].to_csv('inter.txt',
                                sep='\t', index=True, header=True)
        merge.loc[s1_only].to_csv('s1_only.txt',
                                  sep='\t', index=True, header=True)
        merge.loc[s2_only].to_csv('s2_only.txt',
                                  sep='\t', index=True, header=True)
        plt.savefig(labels[0] + '_vs_' + labels[1] , dpi=100)

def draw_venn2(A, B, sets):
    venn = [0]*3
    venn[2] = len(sets["AB"])
    venn[0] = len(sets["A"]) - len(sets["AB"])
    venn[1] = len(sets["B"]) - len(sets["AB"])
    labelA = A + " (" + str(len(sets["A"])) + ")"
    labelB = B + " (" + str(len(sets["B"])) + ")"
    print venn
    venn2(subsets=venn, set_labels=(labelA, labelB))
    plt.show()

def venn(args):
    """
    %prog venn *.benchmark

    Display benchmark results as Venn diagram.
    """
    from matplotlib_venn import venn2

    p = OptionParser(venn.__doc__)
    opts, args, iopts = p.set_image_options(args, figsize="9x9")

    if len(args) < 1:
        sys.exit(not p.print_help())

    bcs = args
    fig = plt.figure(1, (iopts.w, iopts.h))
    root = fig.add_axes([0, 0, 1, 1])

    pad = .02
    ystart = 1
    ywidth = 1. / len(bcs)
    tags = ("Bowers", "YGOB", "Schnable")
    for bc, tag in zip(bcs, tags):
        fp = open(bc)
        data = []
        for row in fp:
            prog, pcounts, tcounts, shared = row.split()
            pcounts = int(pcounts)
            tcounts = int(tcounts)
            shared = int(shared)
            data.append((prog, pcounts, tcounts, shared))
        xstart = 0
        xwidth = 1. / len(data)
        for prog, pcounts, tcounts, shared in data:
            a, b, c = pcounts - shared, tcounts - shared, shared
            ax = fig.add_axes([xstart + pad, ystart - ywidth + pad,
                               xwidth - 2 * pad, ywidth - 2 * pad])
            venn2(subsets=(a, b, c), set_labels=(prog, tag), ax=ax)
            message = "Sn={0} Pu={1}".\
                format(percentage(shared, tcounts, precision=0, mode=-1),
                       percentage(shared, pcounts, precision=0, mode=-1))
            print >> sys.stderr, message
            ax.text(.5, .92, latex(message), ha="center", va="center",
                    transform=ax.transAxes, color='b')
            ax.set_axis_off()
            xstart += xwidth
        ystart -= ywidth

    panel_labels(root, ((.04, .96, "A"), (.04, .96 - ywidth, "B"),
                  (.04, .96 - 2 * ywidth, "C")))
    panel_labels(root, ((.5, .98, "A. thaliana duplicates"),
                        (.5, .98 - ywidth, "14 Yeast genomes"),
                        (.5, .98 - 2 * ywidth, "4 Grass genomes")))
    normalize_axes(root)
    savefig("venn.pdf", dpi=opts.dpi)

def plot_two_set_venn_diagram(Ab, aB, AB, set_labels=None):
    '''
    Plots a two-set venn diagram, A and B.
    
    Inputs:
        Ab: counts in set A but not in B.
        AB: counts in both set A and B (intersection).
        aB: counts in set B but not in A.
        set_labels: a list of length 2 for set A (index 0) and set B (index 1).
    '''
    venn2(subsets=(Ab, Ab, AB), set_labels=set_labels)
    plt.show()

def venn_pair( set1, set2, mask):
    set_a = set()
    for s in linguistic_groups[set1]:
        set_a.update([v.strip() for v in open( mask % ethnicity_code[s]).readlines()])
    set_b = set()
    for s in linguistic_groups[set2]:
        set_b.update([v.strip() for v in open( mask % ethnicity_code[s]).readlines()])

    plt.cla()
    venn2((set_a, set_b), set_labels=(set1, set2))
#    figure.set_size_inches(3,15.4)
#    figure.set_dpi(400)
    plt.savefig("%s_%s_%s.svg" % (set1, set2, 'venn'+mask.split('%s')[1] ))

def main(*fields):
    client = MongoClient()
    db = client['github']
    influences = db['influences']
    sets = map(lambda x: set([i[0] for i in influences.find_one(
        {'field': x})['ranks'][:25]]), fields)

    venn2(sets, fields)
    plt.savefig('images/' +
                fields[0] + '-' + fields[1] +
                '-venn2.png')

    print sets[0].intersection(sets[1])

def main(epistasis_file):
    
    dict_epistasis = {} #list of list of sequences, where each item represents a label 

    with open(epistasis_file) as e:
        lines = e.readlines()
        for l in lines[1:]: #ignore header line
	    tokens = l.split(',')
	    #value consists of Starting Ratio, Ending Ratio, Epistasis, Ending Fitness, # of Mutations, list of InterSeqs, list of InterFits, list of InterRatios
	    if dict_epistasis.get((tokens[2], tokens[0])) is None:
                dict_epistasis[(tokens[0],tokens[2])] = [ float(tokens[1]), float(tokens[3]), float(tokens[5]), tokens[4], len(tokens[6::3]), tokens[6::3], tokens[7::3], tokens[8::3] ]

    neg_epistasis = [0] * 4 
    no_epistasis = [0] * 4
    pos_epistasis = [0] * 4 
    n_functional = [0] * 4
    n_should_be_functional = [0] * 4
    n_total = [0] * 4

    for i in xrange(2,6):
	ind = i-2
        neg_epistasis[ind] = sum([ 1 for key, value in dict_epistasis.items() if value[2] < -0.000005 and value[4] == i ])
        no_epistasis[ind] = sum([ 1 for key, value in dict_epistasis.items() if abs(value[2]) < 0.000005 and value[4] == i ])
        pos_epistasis[ind] = sum([ 1 for key, value in dict_epistasis.items() if value[2] > 0.000005 and value[4] == i ])
	n_functional[ind] = sum([ 1 for key, value in dict_epistasis.items() if value[3] == "CLEAVED" and value[4] == i ])
        n_should_be_functional[ind] = sum([ 1 for key, value in dict_epistasis.items() if all(v == "CLEAVED" for v in value[6]) and value[4] == i ])
	n_total[ind] = float(sum([ 1 for key, value in dict_epistasis.items() if value[4] == i]))

    seq_func = set([ key[1] for key,val in dict_epistasis.items() if val[3] == "CLEAVED" ])
    seq_pred_func = set([ key[1] for key,val in dict_epistasis.items() if all(v == "CLEAVED" for v in val[6]) ]) 
    fig, axarr = pconv.create_ax(1, 1, shx=True, shy=True)
    fig2, axarr2 = pconv.create_ax(1, 1)
    artists = []
    artists.extend(plot_epi(no_epistasis, n_total, axarr[0,0], "No", color="gray"))
    artists.extend(plot_epi(neg_epistasis, n_total, axarr[0,0], "Neg.", bottom=no_epistasis, color="white"))
    artists.extend(plot_epi(pos_epistasis, n_total, axarr[0,0], "Pos.", bottom=[no + neg for no, neg in zip(no_epistasis, neg_epistasis)], color="black"))
    n_func_frac = [ func/total for func, total in zip(n_functional, n_total) ]
    n_pred_frac = [ pred/total for pred, total in zip(n_should_be_functional, n_total) ]
    scatterplot.plot_series(axarr2[0,0], [(range(2,6),n_func_frac,"% Cleaved"),(range(2,6),n_pred_frac,"% Pred Cleaved")], "", "Number of Mutations", "Fraction of Total Cases", size=40, connect_dots=True, alpha=1.0)
    axarr2[0,0].set_ylim([0,4.0])
    fig_venn, axarr_venn = pconv.create_ax(1, 1)

    venn2([seq_func, seq_pred_func], set_labels = ["Cleaved", "Pred Cleaved"], ax=axarr_venn[0,0])

    lgd = axarr[0,0].legend(artists,["No","Neg.","Pos."], loc="center left", bbox_to_anchor=(1.05, 0.5), borderaxespad=0., prop={'size':9}, ncol=1, fancybox=True)


    pconv.save_fig(fig, epistasis_file, "plot", 3, 2.5, tight=False, size=9, extra_artists=lgd)
    pconv.save_fig(fig2, epistasis_file, "pred_v_cl", 5, 5, tight=True, size=10)
    pconv.save_fig(fig_venn, epistasis_file, "venn", 5, 5, tight=True, size=14)

def plot_venn(List_of_sets,
              Set_labels,
              Main = "I forgot to give this plot a name.",
              Out_File = "",
              Custom_overlap_numbers = []):
    """
        Given a list of sets, generate a venn diagram in Out_Dir.
        
        Arguments:
            List_of_sets (two or three only!)
            Set_labels: Label for each circle
            Main: Title of plot
            Out_File: Where should plot be saved? And what should the file be named?
                Parent directory expected to already exist...
                This will overwrite plots if they already exist
            Custom_overlap_numbers: optional. If you want to supply your own 3 overlap sets:
                [# in first, # in second, # in both]
    """
    if not os.path.isdir(os.path.dirname(Out_File)):
        raise ValueError(os.path.dirname(Out_File)+" <--- PATH DOESN'T EXIST")
    
    if len(Custom_overlap_numbers) != 0 and len(Custom_overlap_numbers) != 3:
        raise ValueError("Custom overlap only works for 2 circle venn diagrams at the moment...")
    
    if len(Custom_overlap_numbers) == 3:
        plt.figure()
        venn2(subsets={'10': Custom_overlap_numbers[0], '01': Custom_overlap_numbers[1], '11': Custom_overlap_numbers[2]}, set_labels = Set_labels)
        plt.title(Main)
        plt.savefig(Out_File)
        return
                
    if len(List_of_sets) == 2:
        if len(Set_labels) != 2:
            raise ValueError("Set_labels needs to be the same length as the number of sets...")
        # Default figure dimensions...
        plt.figure()
        venn2(List_of_sets,Set_labels)
        plt.title(Main)
        plt.savefig(Out_File)
        
    elif len(List_of_sets) == 3:
        if len(Set_labels) != 3:
            raise ValueError("Set_labels needs to be the same length as the number of sets...")
        # Default figure dimensions...
        plt.figure()
        venn3(List_of_sets,Set_labels)
        plt.title(Main)
        plt.savefig(Out_File)
    else:
        raise ValueError("List_of_sets needs to be of length 2 or 3.")   

def test_pr_28():
    import matplotlib_venn as mv
    v = mv.venn3((1, 2, 3, 4, 5, 6, 7), subset_label_formatter = None)
    assert v.get_label_by_id('010').get_text() == '2'
    v = mv.venn3((1, 2, 3, 4, 5, 6, 7), subset_label_formatter = lambda x: 'Value: %+0.3f' % (x / 100.0))
    assert v.get_label_by_id('010').get_text() == 'Value: +0.020'
    v = mv.venn2((1, 2, 3), subset_label_formatter = None)
    assert v.get_label_by_id('01').get_text() == '2'
    v = mv.venn2((1, 2, 3), subset_label_formatter = lambda x: 'Value: %+0.3f' % (x / 100.0))
    assert v.get_label_by_id('01').get_text() == 'Value: +0.020'
    
    v = mv.venn3_unweighted((1, 2, 3, 4, 5, 6, 7), subset_label_formatter = lambda x: 'Value: %+0.3f' % (x / 100.0))
    assert v.get_label_by_id('010').get_text() == 'Value: +0.020'
    v = mv.venn2_unweighted((1, 2, 3), subset_label_formatter = lambda x: 'Value: %+0.3f' % (x / 100.0))
    assert v.get_label_by_id('01').get_text() == 'Value: +0.020'

def draw(set1, set2, set3, label1, label2, label3):
  set1 = set(set1)
  set2 = set(set2)
  if label3:
    set3 = set(set3)
    v = venn3([set1,set2, set3], (label1, label2, label3))
    plt.title('Venn diagram for hubs: ' + label1 + "," + label2 +"," + label3, fontsize=20)
  else:
    v = venn2([set1, set2], (label1, label2))
    plt.title('Venn diagram for hubs:' + label1 + "," + label2, fontsize=20)
#   if v.get_label_by_id('110'):
#     plt.annotate(percent_of(set1,set2)+"% of " +label1 , xy=v.get_label_by_id('110').get_position() - np.array([0.15, 0.10]))
#     plt.annotate(percent_of(set2,set1)+"% of " +label2 , xy=v.get_label_by_id('110').get_position() - np.array([0.15, 0.15]))
  if v.get_patch_by_id('100'):
    v.get_patch_by_id('100').set_color("blue")
  if v.get_patch_by_id('010'):
    v.get_patch_by_id('010').set_color("red")
  if v.get_patch_by_id('110'):
    v.get_patch_by_id('110').set_color("purple")
  if label3 and v.get_patch_by_id('001'):
    v.get_patch_by_id('001').set_color("green") 
    if v.get_patch_by_id('111'):
      v.get_patch_by_id('111').set_color("black") 
  gca().set_axis_bgcolor('white')
  gca().set_axis_on()
  plt.show()

def main():
    usage = 'usage: %prog [options] <peaks1_bed> <peaks2_bed> <out_pdf>'
    parser = OptionParser(usage)
    parser.add_option('--l1', dest='label1', default='peaks1', help='Label for peak set 1')
    parser.add_option('--l2', dest='label2', default='peaks2', help='Label for peak set 2')
    (options,args) = parser.parse_args()

    if len(args) != 3:
        parser.error('Must provide two peaks BED files and output PDF')
    else:
        peaks1_bed = args[0]
        peaks2_bed = args[1]
        out_pdf = args[2]

    # count individual
    peaks1_count = count_peaks(peaks1_bed)
    peaks2_count = count_peaks(peaks2_bed)

    # count overlap
    copeaks_count = 0
    p = subprocess.Popen('intersectBed -u -a %s -b %s' % (peaks1_bed, peaks2_bed), stdout=subprocess.PIPE, shell=True)
    for line in p.stdout:
        copeaks_count += 1
    p.communicate()

    plt.figure()
    venn_diag = venn2(subsets=(peaks1_count-copeaks_count, peaks2_count-copeaks_count, copeaks_count), set_labels=[options.label1, options.label2], set_colors=['#e41a1c', '#A1A838'])
    plt.savefig(out_pdf)
    plt.close()

def vdoc_plot(overlap):
	plt.figure(figsize=(13,13), facecolor="white")
	#syntax: set1, set2, set1x2...
	subset_tuple=(5,2,overlap)
	v = venn2(subsets=subset_tuple, set_labels = ('A', 'B', 'C'))

	v.get_patch_by_id('100').set_alpha(0.1)
	v.get_patch_by_id('100').set_color('gray')
	if overlap != 0:
		v.get_patch_by_id('110').set_color('green')
		v.get_patch_by_id('110').set_alpha(0.7)
		v.get_label_by_id('110').set_text('Consciousnes')
	v.get_patch_by_id('010').set_alpha(0.4)
	v.get_label_by_id('100').set_text('Set of all qualia')
	v.get_label_by_id('010').set_text('Set of all concurrent\n mental processes')
	v.get_label_by_id('A').set_text('')
	v.get_label_by_id('B').set_text('')

	c = venn2_circles(subsets=subset_tuple)
	c[0].set_ls('dotted')
	c[1].set_ls('dashed')
	plt.title("Venn Diagram of Consciousnes")

	from matplotlib.transforms import Affine2D
	ax = plt.gca()
	center = [np.mean(ax.get_xlim()), np.mean(ax.get_ylim())]
	t = Affine2D().rotate_deg_around(center[0], center[1], 90) + ax.transData
	for v in ax.patches + ax.texts:
		v.set_transform(t)
	yl = ax.get_ylim()
	plt.ylim(yl[0]-0.2, yl[1]+0.2)

	plt.show()

def plot_venn_ligands(output_path, qu2, mobp2, intersection_mobp):

	""" Plots venn diagram of mOBP vs. Qaio et al. ligands tested in their binding-assays

	Args:

		output_path = "../../data/obp_output/venn_mobp_vs_qaio.png"

		qu2 = see: @lower

		mobp2 = see: @lower

		intersection = see: 

	Usage: 

		plot_venn_ligands("../../data/obp_output/venn_mobp_vs_qaio.png", qu2, mobp2, intersection)

	"""

	plt.figure(figsize=(4,4))
	#v = venn3(subsets=(1, 1, 1, 1, 1, 1, 1), set_labels = ('A', 'B', 'C'))

	v = venn2(subsets = (len(qu2), len(mobp2), len(intersection_mobp)))

	v.get_label_by_id('A').set_text('Exp')
	v.get_label_by_id('B').set_text('MD')


	plt.savefig(output_path)

	plt.close()

def plot_two_set_venn(set1, set2, mycolors=('r', 'g'), 
                      mylabels=None, title='Plot title'):
    '''
    Plot two circle venn diagram
    '''
    # Set matplotlib font size globally
    font = {'family': 'sans',
            'sans-serif': 'Arial',
            'weight': 'bold',
            'size': 25}
    matplotlib.rc('font', **font)
    
    Ab = len(set1 - set2)
    aB = len(set2 - set1)
    AB = len(set1 & set2)
    
    fig = plt.figure()
    # Fill whitespace in the margins by adjusting subplot
    fig.subplots_adjust(bottom=0.10)
    fig.subplots_adjust(left=0.12)
    fig.subplots_adjust(right=0.90)
    fig.subplots_adjust(top=0.90)
    ax = fig.add_subplot(111)
    
    p = venn2(subsets=(Ab, aB, AB), set_colors=mycolors, set_labels=mylabels)
    
    plt.title(title) 
    plt.show()