#! /usr/bin/python


from asset import AssetTaxable, AssetTaxFree, AssetTaxDeferred
from account import AccountTaxable, AccountTaxFree, AccountTaxDeferred
from cash_flow import CashFlow
from portfolio import Portfolio
from copy import deepcopy

class RetirementProjection:
  def __init__(self, periods, retire_period, tax_rate, accounts, savings, spending, verbose=False):
    self.periods       = periods
    self.retire_period = retire_period
    self.tax_rate      = tax_rate
    self.accounts      = deepcopy(accounts)
    self.savings       = deepcopy(savings)
    self.spending      = deepcopy(spending)
 
    self.__orig_accounts = deepcopy(accounts)
    self.__orig_savings  = deepcopy(savings)
    self.__orig_spending = deepcopy(spending)
   
    self.verbose = verbose

    self.portfolio = Portfolio(self.accounts, self.savings)

  def __reset(self):
    self.accounts = deepcopy(self.__orig_accounts)
    self.savings  = deepcopy(self.__orig_savings)
    self.spending = deepcopy(self.__orig_spending)
    self.portfolio = Portfolio(self.accounts, self.savings)

  def logger(self, s):
    if self.verbose: print s

  def run(self):

    taxes_owed = 0
    balances = []
    balances.append(self.portfolio.total())
    for period in range(0, self.periods):
      self.logger("******%s******" % period)
      self.logger("Beg. Total: %s" % self.portfolio.total())


      self.portfolio.apply_savings(period)  

      self.logger("With Savings: %s" %  self.portfolio.total())

      self.portfolio.apply_return()

      self.logger("With Return: %s" % self.portfolio.total())

      spend = self.spending.current_value(period)
      expenses =  spend + taxes_owed

      diff = self.portfolio.total() - expenses
      deficit = diff > 0 and 0 or diff

      self.logger("Retirement Spending: %s" % spend)
      self.logger("Taxes Owed: %s" % taxes_owed)
      self.logger("Total Expenses: %s" % expenses)

      self.portfolio.apply_expenses(expenses)

      self.logger("After Expenses: %s" % self.portfolio.total())
      
      taxable_income = self.portfolio.taxable_income()
      self.logger("Taxable Income: %s" % taxable_income)

      taxes_owed = taxable_income * self.tax_rate

      balances.append(self.portfolio.total()) # XXX need to account for taxes_owed

    self.__reset()
    return balances
    

  def solve_for_savings(self): # this is only for taxable savings. we'll assume qualified accounts are maxed out.
    trial_proj   = deepcopy(self)
    savings_amt  = trial_proj.savings[0].value
    spending_amt = trial_proj.spending.value
    
    if savings_amt == 0: savings_amt = 1000

    results = []
    for i in range(0, 1000):
      print "%s save : %s" % (i, savings_amt)
      
      results = trial_proj.run()
      target = results[-1]
      print "target:", target
      if -1 <= target <= 1 and results[-2] > spending_amt: # within tolerance
        return savings_amt

      if target <= 0:
        savings_amt *= 1.6
      else:
        savings_amt /= 1.4

      trial_proj.savings[0].value = savings_amt

    return savings_amt
    

  def solve_for_spending(self):
    trial_proj = deepcopy(self)
    spending_amt = trial_proj.spending.value
    
    results = []
    for i in range(0, 1000):
      print "%s spend: %s" % (i, spending_amt)
      
      results = trial_proj.run()
      target = results[-1]
      print "target:", target
      if -1 <= target <= 1 and results[-2] > spending_amt: # within tolerance
        return spending_amt

      if target > 0:
        spending_amt *= 1.6
      else:
        spending_amt /= 1.4

      trial_proj.spending.value = spending_amt  

    return spending_amt