import objc, re
from Foundation import *
from AppKit import *

from PyObjCTools import NibClassBuilder
from QIFParser import QIFParser

NibClassBuilder.extractClasses("Drachma", bundle=NSBundle.mainBundle())

def printf(*args):
    print " ".join([str(a) for a in args])
    return

class DummyBudget(object):
    def __init__(self, startDate):
        self._startDate = startDate
    def startDate(self):
        return self._startDate
    def getTheLineItems(self):
        return []

def nextDate(date):
    " Increments the normalized date by one month and returns it "
    nextYear, nextMonth = date
    nextMonth += 1
    if nextMonth > 12:
        nextMonth = 1
        nextYear += 1
    return (nextYear, nextMonth)

def normalize(date):
    """
    Returns something sortable and comparable that represents only the
    significant part.  For example, monthly budgets would return just
    the year and the month.  Currently, we only support monthly
    budgets, so this is hardcoded.
    """
    date = date.dateWithCalendarFormat_timeZone_(None,None)
    return (date.yearOfCommonEra(), date.monthOfYear())

def printable(date):
    " Given the results of normalize(), returns a human readable string "
    userDefaults = NSUserDefaults.standardUserDefaults()
    shortMonthNames = userDefaults.arrayForKey_(NSShortMonthNameArray)
    year, month = date
    return "%s %s" % (shortMonthNames[month-1], year)

class Balancer(NibClassBuilder.AutoBaseClass):

    # -----------------------------------------------------------------
    # Initialization

    def init(self):
        self = super(Balancer, self).init()
        if self is None:
            return None
        # note: nib not loaded yet
        self.a_balances = []
        return self
    
    def awakeFromNib(self):
        self.drachma = self.drachmaController.drachma()
        self.mctx = self.drachma.managedObjectContext()
        pass

    # -----------------------------------------------------------------
    # Queries:
    #
    # Each query against the core data objects is broken up into its
    # own accessor here.  This should allow for easier testing as
    # well.
    
    def sortedBudgets(self):
        entityDesc = NSEntityDescription.entityForName_inManagedObjectContext_(
            u"Budget", self.mctx)
        req = NSFetchRequest.alloc().init()
        req.setEntity_(entityDesc)
        sortdesc = NSSortDescriptor.alloc().initWithKey_ascending_(
            u"startDate", True)
        req.setSortDescriptors_([sortdesc])
        return self.mctx.executeArrayQueryFromFetchRequest_(req)
    
    def categories(self):
        return self.mctx.allObjectsOfEntityType_("Category")
    
    def expensesInCategory(self, cat):
        return cat.valueForKey_("expenses")
    #self.drachma.executeArrayQueryNamed_variables_(
    #        "expensesInCategory", {u"CATEGORY":cat})
    
    def distributeFrom(self, cat):
        return cat.valueForKey_("distributeFrom")

    def distributeInto(self, cat):
        return cat.valueForKey_("distributeInto")

    # -----------------------------------------------------------------
    # Accessors:

    @objc.accessor
    def balances(self):
        """
        A key-value observer compliant accessor for the list of balances.
    
        The list is a list of hashtables that contains information about
        the balance for each category like so:
    
        (Note: all strings are unicode strings)
    
        [
          { 'category':categoryObject,
            'total':total amount of expenses,
            'budgeted':total amount allocated to this category,
            'data':[
              {'date':readable string describing the time period,
               'total':total amount of expenses up to this time period (1),
               'budgeted':total amount budgeted up to this time period (1)
               }* ]
          }*
        ]
    
        *   Indicates more than one item in the list
        (1) Note that these include all previous months
        """
        printf("Fetching a_balances: length=%d" % len(self.a_balances))
        return self.a_balances

    @objc.accessor
    def setBalances_(self, bal): 
        " A key-value observer compliant accessor for the list of balances "
        printf("Setting a_balances to an array of length %s" % len(bal))
        self.a_balances = bal
    
    # -----------------------------------------------------------------
    # Balance manipulation and computation:

    def computeBalances(self):
        " Constructs the data for 'balances' "

        printf("Computing balances...")

        # Determine the list of categories, but exclude those that
        # distribute into other categories.
        categories = [c for c in self.categories()
                      if not self.distributeInto(c)]

        # For each category, accumulate the amount spent in each month
        # separately.  Also, find the date of the first and last expense
        firstDate = NSDate.distantFuture()
        lastDate = NSDate.distantPast()
        expenseCount = 0
        monthlyTotalsByCategory = {}
        for cat in categories:
            # Create a list of expense tuples (e,n) where e is the
            # expense, and n is an integer indicating how many categories
            # the expenses is attributed to.

            printf("  Computing expenses for ", cat.name())
            assert not self.distributeInto(cat)

            ## Those expenses that are in 'cat' belong only to it:
            expenses = [(e, 1) for e in self.expensesInCategory(cat)]
            printf("    %d expenses in this category." % len(expenses))

            ## Those expenses that are in categories that distribute
            ## into 'cat' belong to 'cat' but also others:
            for dcat in self.distributeFrom(cat):
                num = len(self.distributeInto(dcat))
                printf("    Distributed from %s which distributes to %d" % (
                    dcat.name(), num))
                assert num > 0
                dexp = [(e, num) for e in self.expensesInCategory(dcat)]
                expenses += dexp
                printf("      %d expenses from there." % num)

            # Create a dictionary that maps the month to a tuple
            # (amount spent, number of expenses) in this category.
            # Note that not all months will be in here, if no expense
            # occurred.
            monthlyTotals = {}
            for exp, num in expenses:
                expenseCount += 1
                expdate = exp.date()
                if expdate < firstDate: firstDate = expdate
                if expdate > lastDate: lastDate = expdate
                relevantMonth = normalize(expdate)
                assert num > 0
                amount = exp.amount() / num
                try:
                    prevAmount, ctr = monthlyTotals[relevantMonth]
                    ctr += 1
                    prevAmount += amount
                    monthlyTotals[relevantMonth] = prevAmount, ctr
                except KeyError:
                    monthlyTotals[relevantMonth] = amount, 1

            # Sanity check:
            printf("First sanity check for %s:" % cat.name())
            printf("  Number of expenses: %s" % len(expenses))
            totalctr = sum(y for x,y in monthlyTotals.values())
            printf("  Total from monthlyTotals: %s" % totalctr)

            monthlyTotalsByCategory[cat] = monthlyTotals

        printf("firstDate=%s lastDate=%s" % (
            firstDate, lastDate))

        # If there are no categories, we have nothing to list.
        if not expenseCount:
            self.setBalances_([])
            return

        # Now that we have the expenses sorted by category, by month,
        # we combine it with the monthly totals.  We want to compute the
        # total spent and budgeted on a category overall, as well as
        # running totals for each month.
        categoriesData = {}
        budgets = self.sortedBudgets()
        if not budgets: budgets = [DummyBudget(firstDate)]
        for idx in range(len(budgets)):
            budget = budgets[idx]
    
            # Find the start and end date for this budget.  We know
            # that these always correspond to midnight on the first of
            # some month; if this is the last budget, make the end
            # date be the beginning of next month.
            startDate = normalize(budget.startDate())
            if idx == len(budgets)-1:
                endDate = nextDate(normalize(lastDate))
            else:
                endDate = normalize(budgets[idx+1].startDate())

            printf("Budget from %s to %s " % (startDate, endDate))

            # If there are no expenses in this budget's time frame, stop
            if startDate > normalize(lastDate):
                break
            printf("  normalize(lastDate)=",normalize(lastDate))
    
            # For each category...
            for cat in categories:

                printf("  Category: %s" % cat.name())
                
                # Initialize the data for this category if this is the
                # first time we see it, otherwise re-use it.
                try:
                    catData = categoriesData[cat]
                except KeyError:
                    catData = {
                        u'category':cat,
                        u'total':NSDecimalNumber.zero(),
                        u'budgeted':NSDecimalNumber.zero(),
                        u'balance':NSDecimalNumber.zero(),
                        u'count':0,
                        u'data':[]}
                    categoriesData[cat] = catData
                printf("    Category data initialized")

                # Get the monthly totals just for this category.
                monthlyTotals = monthlyTotalsByCategory[cat]
                printf("    Monthly totals loaded")
                
                # Find the line item in the budget (if any)
                for lineItem in budget.getTheLineItems():
                    if lineItem.category() == cat:
                        printf("    Line item located")
                        budgetAmount = lineItem.amount()
                        break
                else:
                    budgetAmount = NSDecimalNumber.zero()
                printf("    Budget Amount == %s" % budgetAmount.description())
                
                # Fill in the data for each month, and compute a
                # hashtable for each one with the date and running
                # totals so far
                date = startDate
                totalSoFar = catData[u'total']
                budgetedSoFar = catData[u'budgeted']
                catDataList = catData[u'data']
                counterSoFar = catData[u'count']
                while date < endDate:
                    if date in monthlyTotals:
                        thisMonth, thisMonthCounter = monthlyTotals[date]
                        totalSoFar += thisMonth
                    else:
                        thisMonth = NSDecimalNumber.zero()
                        thisMonthCounter = 0
                    counterSoFar += thisMonthCounter
                    budgetedSoFar += budgetAmount
                    printf("    Date=%s, count=%s total=%s, budget=%s" % (
                        date, counterSoFar, totalSoFar, budgetedSoFar))
                    catDataList.append({
                        u'date':printable(date),
                        u'amount':thisMonth,
                        u'count':thisMonthCounter,
                        u'budgeted':budgetAmount,
                        u'balance':budgetAmount-thisMonth,
                        u'totalSoFar':totalSoFar,
                        u'budgetedSoFar':budgetedSoFar,
                        u'balanceSoFar':budgetedSoFar-totalSoFar})
                    date = nextDate(date)
                catData[u'total'] = totalSoFar 
                catData[u'budgeted'] = budgetedSoFar
                catData[u'balance'] = budgetedSoFar-totalSoFar
                catData[u'count'] = counterSoFar

        # Perform a sanity check
        for cat in categories:
            printf("Category=%s" % cat.name())
            catData = categoriesData[cat]
            printf("  count=%s" % catData[u'count'])
            printf("  total=%s" % catData[u'total'])
            expenses = self.expensesInCategory(cat)
            printf("  recomputed count=%s" % len(expenses))
            expensesTotal = sum(exp.amount() for exp in expenses)
            printf("  recomputed total=%s" % expensesTotal)
    
        self.setBalances_(categoriesData.values())    
    
    pass