#!/usr/bin/env python

try:
    from Queue import Queue, Empty as QueueEmpty
except ImportError:
    from queue import Queue, Empty as QueueEmpty
import argparse
try:
    from itertools import izip_longest, product
except ImportError:
    from itertools import zip_longest as izip_longest, product
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np
import threading
import time
import sys


class XYGenerator(object):
    def __init__(self, first_is_x=False, xy=False):
        self.first_is_x = first_is_x
        self.xy = xy
        self._data = None
        self._idx = 0

    def n_pairs(self, data):
        if self.first_is_x:
            return data.shape[1] - 1
        if self.xy:
            return data.shape[1] / 2
        return data.shape[1]

    def __call__(self, data):
        self._data = data
        self._idx = 0
        return self

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        i = j = None

        if self.first_is_x:
            i = 0
            j = self._idx
            self._idx += 1

        elif self.xy:
            i = self._idx
            j = self._idx + 1
            self._idx += 2

        else:
            j = self._idx
            self._idx += 1

        if j >= self._data.shape[1]:
            raise StopIteration()

        idx_mask = np.logical_not(np.isnan(self._data[:, j]))
        idxs = np.arange(len(self._data))[idx_mask]
        if i is None:
            return (idxs, self._data[idxs, j]), j
        else:
            return (self._data[idxs, i], self._data[idxs, j]), j


def last(iterable):
    for x in iterable:
        pass
    return x


def create_slice(s):
    def none_or_int(x):
        try:
            return int(x.strip())
        except ValueError:
            return None

    return slice(*list(map(none_or_int, s.split(":"))))


def create_filter(s):
    def inner(x):
        coeff = np.array(list(map(float, s.split(","))))
        coeff /= coeff.sum()
        return np.hstack([
            np.convolve(x[:, i], coeff, mode="valid")[:, np.newaxis]
            for i in range(x.shape[1])
        ])
    return inner


def plot(
    fig, ax,
    data, labels, xy_pairs,
    subplots, logscale,
    title, xlabel, ylabel,
    ylim
):
    lines = []
    if subplots:
        N, _ = ax.shape
        axes = product(range(N), range(N))
        for xy_label, ij in izip_longest(xy_pairs(data), axes):
            if xy_label:
                xy, label_idx = xy_label
                lines.extend(ax[ij].plot(*xy))
                if labels[label_idx] is not None:
                    ax[ij].set_title(labels[label_idx])
                if logscale:
                    ax[ij].set_yscale("log")
                ax[ij].set_xlabel(xlabel)
                ax[ij].set_ylabel(ylabel)
                ax[ij].set_ylim(ylim)
                ax[ij].grid()
            else:
                ax[ij].set_axis_off()
        fig.suptitle(title)
    else:
        for xy, label_idx in xy_pairs(data):
            lines.extend(ax.plot(*xy, label=labels[label_idx]))
        if logscale:
            ax.set_yscale("log")
        ax.set_title(title)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.set_ylim(ylim)
        ax.grid()
        plt.legend()

    return lines


def update_plot(data, lines, ax, xy_pairs, filter, slice, ylim):
    N = ax.shape[1] if isinstance(ax, np.ndarray) else None
    filtered = filter(data)
    xmin = ymin = float("inf")
    xmax = ymax = -float("inf")
    if len(filtered[slice]) <= 1:
        return lines, ax
    for i, (l, ((x, y), _)) in enumerate(zip(lines, xy_pairs(filtered))):
        l.set_xdata(x[slice])
        l.set_ydata(y[slice])
        if N is not None:
            r, c = divmod(i, N)
            if not ylim:
                ax[r, c].set_ylim((y.min(), y.max()))
            ax[r, c].set_xlim((x.min(), x.max()))
        else:
            xmin, xmax = min(x.min(), xmin), max(x.max(), xmax)
            ymin, ymax = min(y.min(), ymin), max(y.max(), ymax)
    if N is None:
        if not ylim:
            ax.set_ylim(ymin, ymax)
        ax.set_xlim(xmin, xmax)

    return lines, ax


def read_lines():
    def reading_thread(q):
        while True:
            l = sys.stdin.readline()
            if l == "":
                break
            q.put(l)

    q = Queue()
    t = threading.Thread(target=reading_thread, args=(q,))
    t.daemon = True
    t.start()

    return t, q


def data_gen(data, delimiter, stop_at_eof=False):
    t, q = read_lines()

    while True:
        try:
            line = q.get_nowait()
            data = np.vstack([
                data,
                np.array(list(map(float, line.split(delimiter)))).reshape(1, -1)
            ])
        except QueueEmpty:
            yield data
            if not t.is_alive() and stop_at_eof:
                break


def main(argv):
    parser = argparse.ArgumentParser(
        description="Plot data with lines from stdin"
    )

    parser.add_argument(
        "--first_is_x",
        action="store_true",
        help="Use the first number for the x axis"
    )
    parser.add_argument(
        "--xy",
        action="store_true",
        help="Use two consecutive columns as xy pairs"
    )
    parser.add_argument(
        "--delimiter", "-d",
        help="Split on this delimiter"
    )
    parser.add_argument(
        "--subplots",
        action="store_true",
        help="Plot the lines into a NxN subplot grid"
    )
    parser.add_argument(
        "--logscale",
        action="store_true",
        help="Use logscale for the y axis"
    )
    parser.add_argument(
        "--slice",
        type=create_slice,
        help="Slice the dataset",
        default=slice(None)
    )
    parser.add_argument(
        "--filter",
        type=create_filter,
        default=lambda x: x,
        help="Convolve data with a 1D filter with those coefficients"
    )
    parser.add_argument(
        "--title",
        default="",
        help="Set the title in the case of single plot"
    )
    parser.add_argument(
        "--xlabel",
        default="",
        help="Set the xlabel in the case of single plot"
    )
    parser.add_argument(
        "--ylabel",
        default="",
        help="Set the xlabel in the case of single plot"
    )
    parser.add_argument(
        "--ylim",
        type=lambda x: list(map(float, x.split(","))),
        default=None,
        help="Specify the limits of the y axis"
    )
    parser.add_argument(
        "--non_interactive",
        dest="interactive",
        action="store_false",
        help="Don't update the plot with new points read from stdin"
    )
    parser.add_argument(
        "--figsize",
        type=lambda x: list(map(float, x.split(","))),
        default=None,
        help="Specify the size of the figure"
    )

    args = parser.parse_args(argv)

    # Try to figure out get the labels and the first data
    firstline = sys.stdin.readline()
    try:
        data = np.array(
            list(map(float, firstline.split(args.delimiter)))
        ).reshape(1, -1)
        labels = [None]*data.shape[1]
    except ValueError:
        labels = firstline.split(args.delimiter)
        firstline = sys.stdin.readline()
        data = np.array(
            list(map(float, firstline.split(args.delimiter)))
        ).reshape(1, -1)

    # Create a XYGenerator to return the xy pairs for plotting
    xy_pairs = XYGenerator(args.first_is_x, args.xy)

    # Create the figure and the axes
    kwargs = dict()
    if args.figsize:
        kwargs["figsize"] = tuple(args.figsize)
    if args.subplots:
        N = int(np.ceil(np.sqrt(xy_pairs.n_pairs(data))))
        fig, ax = plt.subplots(N, N, **kwargs)
    else:
        fig, ax = plt.subplots(**kwargs)

    if args.interactive:
        lines = plot(
            fig, ax,
            data, labels, xy_pairs,
            args.subplots, args.logscale,
            args.title, args.xlabel, args.ylabel, args.ylim
        )
        anim = FuncAnimation(
            fig, update_plot,
            frames=data_gen(data, args.delimiter),
            fargs=(lines, ax, xy_pairs, args.filter, args.slice, args.ylim),
            interval=2000
        )
    else:
        data = last(data_gen(data, args.delimiter, stop_at_eof=True))
        lines = plot(
            fig, ax,
            args.filter(data)[args.slice], labels, xy_pairs,
            args.subplots, args.logscale,
            args.title, args.xlabel, args.ylabel, args.ylim
        )
    plt.show()


if __name__ == "__main__":
    main(sys.argv[1:])
