#!/usr/bin/env python

from pylab import *
from string import join
from random import shuffle

def stacked_filled_plot(x,ys,**kw):
  try:
    sp = kw['subplot']
  except:
    sp = False
  try:
    colors = kw['colors']
  except:
    colors = ["red",
              "blue",
              "green",
              "cyan",
              "magenta",
              "yellow",
              "brown",
              "silver",
              "purple",
              "lightyellow",
              "black",
              "ivory",
              "pink",
              "orange",
              "gray",
              "teal"]
 
  if sp:
    subplot(sp)
  ax = array(x)
  x2 = concatenate((ax,ax[::-1]))
  y_top = zeros(len(ax))
  paths = []    # a vestage of the days when I tried to use a true legend
  try:
    names = kw['names']
  except:
    names = False

  name_color = {}
  for i in range(len(ys)):
    new_y_top = y_top + array(ys[i])
    p = fill(x2,concatenate((new_y_top,y_top[::-1])),
             colors[i % len(colors)])
    paths.append(p)

    if names:
      name_color[names[i]] = colors[i % len(colors)]
    y_top = new_y_top

  if names:
    names.reverse()
    if sp:
      subplot(sp+1)
      axis("off")
    cols = 2
    rows = (len(names) + (cols - 1)) / cols

    cells = []
    cell_colors = []
    cell_index = 0
    for i in range(rows):
      row = []
      row_colors = []
      for j in range(cols):
        if cell_index < len(names):
          row.append(names[cell_index])
          row_colors.append(name_color[names[cell_index]])
        else:
          row.append("")
          row_colors.append("w")
        cell_index += 1
      cells.append(row)
      cell_colors.append(row_colors)

      table(cellText=cells,cellColours=cell_colors, loc='center')

  return name_color


if __name__ == "__main__":
  x = arange(10)
  y = rand(10,10)
  names = [" very very very very long series name %d" % (i) for i in x]
  stacked_filled_plot(x,y,names=names,subplot=211)
  show()