import time
import os
import pybedtools
from pybedtools.contrib import plotting
from matplotlib import pyplot as plt
import matplotlib
          #res           b          a
colors = ['#A1D99B', '#FDBB84', '#9ECAE1']

font = {'family' : 'sans',
        'weight' : 'normal',
        'size'   : 12}

matplotlib.rc('font', **font)

def plot_a_b_tool(a, b, method, title, labelA, labelB, **kwargs):
    """
    Use for BEDTools programs that use -a and -b input arguments.  Filenames
    `a` and `b` are used for `method` of the BedTool class, and `kwargs` are
    sent to that method.

    The result is a plot of `a`, `b`, and the result, with the commandline
    argument as the plot title.
    """
    a = pybedtools.BedTool(a, from_string=True)
    b = pybedtools.BedTool(b, from_string=True)
    kwargs['b'] = b
    result = getattr(a, method)(**kwargs)

    fig = plt.figure(figsize=(8, 2))
    ax = fig.add_subplot(111)
    ybase = 0
    yheight = 0.1
    ylabels = []
    yticks = []
    for color, bt, label in zip(
            colors,
            [result, b, a],
            ['Result', labelA, labelB]):
        ylabels.append(label)
        track = plotting.Track(
                bt, visibility='squish', alpha=0.9, ybase=ybase, yheight=yheight,
                color=color)
        yticks.append(track.midpoint)
        ybase = track.ymax + 0.1
        ax.add_collection(track)
    ax.set_yticklabels(ylabels)
    ax.set_yticks(yticks)
    #ax.set_title(' '.join([os.path.basename(i) for i in result._cmds]))
    ax.set_title(title)
    ax.axis('tight')
    fig.subplots_adjust(top=0.8, bottom=0.15)
    return ax