"""
This module provides plotting filters.
"""
import itertools

from penchy.jobs.elements import Filter
from penchy.jobs.typecheck import Types
from penchy.util import default, average


class Plot(Filter):
    def __init__(self, filename, title="", xlabel="",
                 ylabel=""):
        super(Plot, self).__init__()
        self.filename = filename
        self.title = title
        self.xlabel = xlabel
        self.ylabel = ylabel


class BarPlot(Plot):
    """
    A simple and static barplot.  Mainly for testing.
    """
    inputs = Types(('x', list, str), ('y', list))

    #FIXME: Better solution to handle args without repeating
    def __init__(self, colors, zlabels, width=0.2, **kwarg):
        super(BarPlot, self).__init__(**kwarg)
        self.width = width
        self.colors = colors
        self.zlabels = zlabels

    def _run(self, **kwargs):
        import numpy as np
        import matplotlib.pyplot as plt
        xs = kwargs['x']
        yss = kwargs['y']
        ind = np.arange(len(xs))
        fig = plt.figure()
        plot = fig.add_subplot(1, 1, 1)
        bars, names, rects = [], [], []
        for i, ys, c, zlabel in zip(itertools.count(), zip(*yss), self.colors, self.zlabels):
            rects.append(plot.bar(ind + self.width * i, ys, self.width, color=c))
            bars.append(rects[i][0])
            names.append(zlabel)
        plot.set_xlabel(self.xlabel)
        plot.set_ylabel(self.ylabel)
        plot.set_title(self.title)
        plot.set_xticks(ind + self.width)
        plot.set_xticklabels(xs)
        plot.legend(bars, names)
        plt.savefig(self.filename)


class ScatterPlot(Plot):
    """
    Scatterplot
    """
    inputs = Types(('x', list, (int, float)),
                   ('y', list, (int, float)),
                   ('z', list, str))

    #FIXME: Better solution to handle args without repeating
    #TODO: allow to draw circles and squares
    def __init__(self, **kwarg):
        super(ScatterPlot, self).__init__(**kwarg)

    def _run(self, **kwargs):
        import numpy as np
        import matplotlib.pyplot as plt
        xs = kwargs['x']
        ys = kwargs['y']
        zs = kwargs['z']
        fig = plt.figure()
        plot = fig.add_subplot(1, 1, 1)

        for x, y, z in zip(xs, ys, zs):
            plot.text(x, y, z, rotation=-45)
        plot.plot(xs, ys, 'o')

        plot.set_xlabel(self.xlabel)
        plot.set_ylabel(self.ylabel)
        plot.set_title(self.title)
        plt.savefig(self.filename)

@@ -478,57 +478,15 @@
     
-A simple and static barp
+BarP
 lot
-.  Mainly for testing.
 %0A   
@@ -641,16 +641,34 @@
 zlabels,
+ horizontal=False,
  width=0
@@ -814,16 +814,53 @@
  zlabels
+%0A        self.horizontal = horizontal
 %0A%0A    de
@@ -1141,16 +1141,34 @@
  %5B%5D, %5B%5D%0A
+        print yss%0A
         
@@ -1249,24 +1249,186 @@
 f.zlabels):%0A
+            print ys%0A            if self.horizontal:%0A                rects.append(plot.barh(ind + self.width * i, ys, self.width, color=c))%0A            else:%0A    
             
@@ -1489,32 +1489,32 @@
 idth, color=c))%0A
-
             bars
@@ -1668,32 +1668,161 @@
 tle(self.title)%0A
+        if self.horizontal:%0A            plot.set_yticks(ind + self.width)%0A            plot.set_yticklabels(xs)%0A        else:%0A    
         plot.set
@@ -1839,32 +1839,36 @@
 d + self.width)%0A
+    
         plot.set

allow horizontal barplots