#!/usr/bin/env python
"""
Given 3 files, creates a 3-way Venn diagram of intersections using the Google \
Chart API; see :mod:`pybedtools.contrib.venn_maker` for more flexibility.

The values in the diagram assume:

    * unstranded intersections
    * no features that are nested inside larger features
"""

import argparse
import sys
import pybedtools
from six.moves import urllib


def venn_gchart(a, b, c=None, colors=None, labels=None, size='300x300'):
    """
    a, b, and c are filenames to BED-like files.

    *colors* is a list of 3 hex colors

    *labels* is a list of 3 labels

    *outfn* is the output PNG you want to create.

    *size* is the size in pixels for the PNG
    """
    a = pybedtools.BedTool(a)
    b = pybedtools.BedTool(b)
    if c:
        c = pybedtools.BedTool(c)

    # The order of values is meaningful to the API, see
    # http://code.google.com/apis/chart/docs/gallery/venn_charts.html
    if c:
        vals = [len(a),
                len(b),
                len(c),
                len(a + b),
                len(a + c),
                len(b + c),
                len(a + b + c)]
    else:
        # insert 0 for size of 3rd circle.
        vals = [len(a), len(b), 0, len(a + b)]
        labels = labels[:2]
    # API doesn't seem to like large numbers, so get fractions instead, then
    # join make a comma-separated list of values.
    mx = float(max(vals))
    vals = [round(i / mx, 4) for i in vals]
    valstr = ','.join(map(str, vals))

    data = {'cht': 'v',
            'chs': size,
            'chd': 't:' + valstr}

    # Add the optional data, if specified
    if labels:
        data['chdl'] = '|'.join(labels)
    if colors:
        data['chco'] = ','.join(colors)
    return data


def gchart(data, outfn='out.png'):
    """
    Sends data to Google Chart API
    """
    data = urllib.parse.urlencode(data)
    binary_data = data.encode('UTF-8')

    url = 'https://chart.googleapis.com/chart?'

    # Request and get the PNG
    req = urllib.request.Request(url, binary_data)
    print(url + data)
    response = urllib.request.urlopen(req)
    f = open(outfn, 'w')
    f.write(response.read())
    f.close()


def main():
    """Create a 3-way Venn diagram using Google Charts API
    """

    op = argparse.ArgumentParser(description=__doc__, prog=sys.argv[0])
    op.add_argument('-a', help='File to use for the left-most circle')
    op.add_argument('-b', help='File to use for the right-most circle')
    op.add_argument('-c', help='File to use for the bottom circle')
    op.add_argument('--colors', help='Optional comma-separated list of hex '
                    'colors for circles a, b, and c.  E.g. %(default)s',
                     default='00FF00,FF0000,0000FF')
    op.add_argument('--labels',
            help='Optional comma-separated list of labels for a, b, and c',
                    default='a,b,c')
    op.add_argument('--size', default='300x300',
                  help='Optional size of PNG, in pixels.  Default is '
                  '"%(default)s"')
    op.add_argument('-o', default='out.png',
                  help='Output file to save as, in PNG format')
    op.add_argument('--test', action='store_true',
            help='run test, overriding all other options.')
    options = op.parse_args()

    reqd_args = ['a', 'b']
    if not options.test:
        for ra in reqd_args:
            if not getattr(options, ra):
                op.print_help()
                sys.stderr.write('Missing required arg "%s"\n' % ra)
                sys.exit(1)

    if options.test:
        # Example data
        pybedtools.bedtool.random.seed(1)
        a = pybedtools.example_bedtool('rmsk.hg18.chr21.small.bed')
        b = pybedtools.example_bedtool('venn.b.bed')
        c = pybedtools.example_bedtool('venn.c.bed')
        options.a = a.fn
        options.b = b.fn
        options.c = c.fn
        options.colors = '00FF00,FF0000,0000FF'
        options.o = 'out.png'
        options.labels = 'a,b,c'

    data = venn_gchart(a=options.a, b=options.b, c=options.c,
                       colors=options.colors.split(','),
                       labels=options.labels.split(','),
                       size=options.size)
    gchart(data, outfn=options.o)

if __name__ == "__main__":
    import doctest
    if doctest.testmod(optionflags=doctest.ELLIPSIS).failed == 0:
        main()