#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
plot data with `matplotlib`

See also :
- http://stackoverflow.com/questions/7534453/matplotlib-does-not-show-my-drawings-although-i-call-pyplot-show ;
- http://bokeh.pydata.org/ ;
- http://mpld3.github.io/
"""

# imports
import argparse
import csv
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import numpy as np
import os
import subprocess
import sys
import tempfile
import time
from .utils import choose_program
from which import which
from StringIO import StringIO

# module globals
__all__ = ['Displayer', 'Plot', 'PlotParser', 'read', 'main']
string = (str, unicode)


class Displayer(object):
    """image displayer"""
    viewers = ('feh', 'qiv', 'sxiv')

    def __init__(self, viewers=None):
        if viewers is None:
            viewers = self.viewers
        self.viewers = viewers

    def viewer(self):
        """returns path to primary viewer"""
        for viewer in self.viewers:
            path = which(viewer)
            if path:
                return path

    def __call__(self, image):
        """display an image and exit"""
        viewer = self.viewer()
        if viewer:
            args = [viewer, image]
            subprocess.check_output(args)


class Plot(object):
    """plotting class"""
    def __init__(self, title=None, xlabel=None, ylabel=None):
        self.title = title
        self.xlabel = xlabel
        self.ylabel = ylabel
        self._x = None
        self._y = []
        self.marker = []

    def x(self, data, label=None):
        self._x = data
        if label is not None:
            self.xlabel = xlabel

    def y(self, data, label=None, marker='.'):
        self._y.append(data)
        self.marker.append(marker)
        if label is not None:
            self.ylabel = label

    def __call__(self, output):
        assert self._y

        if self.title:
            plt.title(self.title)
        if self.xlabel:
            plt.xlabel(self.xlabel)
        if self.ylabel:
            plt.ylabel(self.ylabel)
        if self._x:
            args = sum([[self._x, self._y[i], self.marker[i]] for i in range(len(self._y))], [])
            plt.plot(*args)
        else:
            plt.plot(*self._y)
        plt.show()
        plt.savefig(output)
        print ("{}->saved to '{}'".format(self.title or '', output))


def read(f):
    """
    Read from file ``f``
    Accepts CSV and space-delimited files
    """

    retval = None
    for line in f:
        line = line.strip()
        if ',' in line:
            buffer = StringIO()
            buffer.write(line)
            buffer.seek(0)
            row = list(csv.reader(buffer))[0]
        else:
            row = line.split()
        row = [float(i) for i in row]
        if retval is None:
            retval = [[i] for i in row]
        else:
            for index, value in enumerate(row):
                retval[index].append(value)

    return retval


class PlotParser(argparse.ArgumentParser):
    """CLI option parser for the plotter"""

    def __init__(self, *args, **kwargs):
        kwargs.setdefault('description', __doc__)
        argparse.ArgumentParser.__init__(self, *args, **kwargs)
        self.add_argument('--info', dest='info',
                            help="display info and exit")
        self.add_argument('input',
                          type=argparse.FileType('r'), nargs='*',
                          default=(sys.stdin,),
                          help='input file(s), or read from stdin if ommitted')
        self.add_argument('-o', '--output', dest='output',
                          help="file name to output to")
        self.add_argument('-s', '--scatter', dest='scatter',
                          action='store_true', default=False,
                          help="scatter plot")
        self.options = None

    def parse_args(self, *args, **kwargs):
        options = argparse.ArgumentParser.parse_args(self, *args, **kwargs)
        self.validate(options)
        self.options = options
        return options

    def validate(self, options):
        """validate options"""


def main(args=sys.argv[1:]):
    """CLI"""

    # parse command line
    parser = PlotParser()
    options = parser.parse_args(args)

    plot_fcn = plt.scatter if options.scatter else plt.plot

    # read data
    all_data = [read(f) for f in options.input]

    # color map
    # http://stackoverflow.com/questions/12236566/setting-different-color-for-each-series-in-scatter-plot-on-matplotlib
    n_col = sum([(len(data)-1) or 1 for data in all_data])
    colors = iter(cm.rainbow(np.linspace(0, 1, n_col)))

    for data in all_data:
        # plot it
        if len(data) == 1:
            plot_fcn(*data, marker='.', color=next(colors))
        else:
            for i in range(1, len(data)):
                plot_fcn(data[0], data[i], label=str(i), marker='.', color=next(colors))

    plt.show()

    # display
    output = options.output or tempfile.mktemp(suffix='.png')
    plt.savefig(output)
    Displayer()(output)

if __name__ == '__main__':
    main()