"""
Generate various plots
"""
import matplotlib.pyplot as plt
import numpy as np

from bluechips import model
from bluechips.model import meta, User, Expenditure

def balance_plot(filename='balance', dates=None):
    """
    Create a plot showing the net balance for each user over the
    course of time.

    dates is a tuple of datetime objects.
    """
    users = meta.Session.query(User).all()
    expenses = meta.Session.query(Expenditure).order_by(Expenditure.date)

    totals = dict()
    dates = [e.date for e in expenses]
    for u in users:
        totals[u.name] = np.cumsum([e.share(u) for e in expenses])
    
    fig = plt.figure()
    ax = fig.add_subplot(1,1,1)
    lines = dict()
    for u,y in totals.items():
        lines[u] = plt.plot(dates,totals[u], label=u)
        
    ax.legend(loc='best')
    fig.savefig(filename)