"""
Generate various plots
"""
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import numpy as np
import datetime

from cStringIO import StringIO
from pylons import request, response

from bluechips.lib.base import *
import bluechips.lib.totals
from bluechips import model
from bluechips.model import meta, orm, User, Expenditure, Split, Transfer

import time

def get_share(transaction, user):
    if type(transaction) == Expenditure:
        amount = transaction.share(user)
        if transaction.spender == user:
            amount -= transaction.amount
        return amount
    else:
        if user == transaction.debtor:
            return -transaction.amount
        elif user == transaction.creditor:
            return transaction.amount
        else:
            return model.types.Currency(0)

class GraphController(BaseController):
    def balance(self):
        """
        Create a plot showing the net balance for each user over the
        course of time.
        """
        t0 = time.time()
        users = meta.Session.query(User).filter(User.resident==True).all()
        expenditures = meta.Session.query(Expenditure).\
                       options(orm.eagerload_all(Expenditure.splits, Split.user)).\
                       all()
        transfers = meta.Session.query(Transfer).\
                    options(orm.eagerload(Transfer.debtor, Transfer.creditor)).\
                    all()

        transactions = expenditures + transfers
        transactions.sort(key=lambda x: x.date)

        t1 = time.time()
        totals = dict()
        dates = [t.date for t in transactions]
        for u in users:
            totals[u.name] = np.cumsum([get_share(t,u) for t in transactions])

        t2 = time.time()
        fig = plt.figure()
        canvas = FigureCanvas(fig)
        ax = fig.add_subplot(1,1,1)
        ax.set_color_cycle(plt.cm.jet(np.linspace(0,1,len(users))))
        lines = dict()
        ymax = 0
        for u,y in totals.items():
            lines[u] = plt.step(dates,totals[u], label=u, lw=2)

        ax.legend(loc='best')
        ax.grid(True)
        plt.axis('tight')
        ax.set_xlim(xmin=datetime.date.today()-datetime.timedelta(days=365))
        ax.set_ylim(ymin=0)
        fig.autofmt_xdate(rotation=30)
        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter("%h"))
        ax.set_yticklabels(map(lambda x: '%s$%i'%('-' if x<0 else '',abs(x)//100),
                               ax.yaxis.get_majorticklocs()))

        t3 = time.time()
        s = StringIO()
        canvas.print_figure(s)
        response.headers['Content-Type'] = 'image/png'
        t4 = time.time()
#        print t1-t0, t2-t1, t3-t2, t4-t3
        return s.getvalue()

    def pie(self):
        ruser = request.environ['user']

        debts = []
        debtors = []
        credits = []
        creditors = []
        explode_debt = []
        explode_credit = []
        eps = model.types.Currency("2.00")
        net = bluechips.lib.totals.debts().items()
        net.sort(key=lambda x: x[1])
        for user, amount in net:
            if amount > eps:
                debtors.append(user.name)
                debts.append(amount)
                explode_debt.append(0.1 if user == ruser else 0)
            elif amount < -eps:
                creditors.append(user.name)
                credits.append(-amount)
                explode_credit.append(0.1 if user == ruser else 0)

        fig = plt.figure(figsize=(6,3))
        canvas = FigureCanvas(fig)

        ax1 = fig.add_axes([0.05,0.1,0.4,0.8])
        ax1.pie(debts,
                explode=explode_debt,
                labels=debtors,
                colors=plt.cm.YlOrRd(np.linspace(0.05,0.95,len(debtors))))
        ax1.set_title('Debtors')

        ax2 = fig.add_axes([0.55,0.1,0.4,0.8])
        ax2.pie(credits,
                explode=explode_credit,
                labels=creditors,
                colors=plt.cm.Blues(np.linspace(0.05,0.95,len(creditors))))
        ax2.set_title('Creditors')

#        ax.legend(loc='best')
#        ax.grid(True)
#        plt.axis('tight')
#        ax.set_xlim(xmin=datetime.date.today()-datetime.timedelta(days=365))
#        ax.set_ylim(ymin=0)
#        fig.autofmt_xdate(rotation=30)
#        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter("%h"))
#        ax.set_yticklabels(map(lambda x: '%s$%i'%('-' if x<0 else '',abs(x)//100),
#                               ax.yaxis.get_majorticklocs()))

        s = StringIO()
        canvas.print_figure(s, dpi=72)
        response.headers['Content-Type'] = 'image/png'
        return s.getvalue()