#! /usr/bin/python

from asset import AssetTaxable, AssetTaxFree, AssetTaxDeferred
from utility import lineno
from pprint import pprint

class Account:

  def __init__(self, name):
    self.name = name

    self.withdrawal_penalty = 0
    self.capital_income     = 0
    self.ordinary_income    = 0
    self.assets = []

  def value(self):
    total = 0
    for asset in self.assets:
      total += asset.value
    return total

  def apply_return(self):
    for asset in self.assets:
      asset.apply_return()

  # it's the same for each account since it does not result in any taxable events.
  # rebalancing will come after this. -- note not very tax efficient.
  # maybe to reduce taxes we can apply the savings according to allocation needs.
  def apply_savings(self, savings):
    total = self.value()
    if total > 0:
      for asset in self.assets:
        asset.adjust_value(savings * (float(asset.value) / total))
    else:
      nassets = len(self.assets)
      for asset in self.assets:
        asset.adjust_value(float(savings) / nassets)

  def get_capital_income(self):
    value = self.capital_income
    self.capital_income = 0
    return value

  def get_withdrawal_penalty(self):
    value = self.withdrawal_penalty
    self.withdrawal_penalty = 0
    return value

class AccountTaxable(Account):
  def apply_expenses(self, expenses):
    expenses = -abs(expenses) # this way we don't care if an expense is neg. or pos. ... it's always neg.
    if abs(expenses) == 0: return 0
    
    for asset in self.assets:
      expenses = asset.adjust_value(expenses)
      self.capital_income += abs(asset.taxable_income) 
      if abs(expenses) == 0: 
        break
      
    return expenses

class AccountTaxDeferred(Account):
  def apply_expenses(self, expenses, penalty_applies=False):
    expenses = -abs(expenses) # this way we don't care if an expense is neg. or pos. ... it's always neg.
    if abs(expenses) == 0: return 0
    
    for asset in self.assets:
      orig_value = asset.value
      expenses = asset.adjust_value(expenses)
      self.ordinary_income += abs(asset.taxable_income)
      if penalty_applies: 
        self.withdrawal_penalty += 0.2 * (orig_value - float(asset.value))
      if abs(expenses) == 0: break

    return expenses

class AccountTaxFree(Account):
  def apply_expenses(self, expenses, penalty_applies=False):
    expenses = -abs(expenses) # this way we don't care if an expense is neg. or pos. ... it's always neg.
    if abs(expenses) == 0: return 0
    
    for asset in self.assets:
      expenses = asset.adjust_value(expenses)
      if penalty_applies:
        self.withdrawal_penalty += 0.2 * (orig_value - float(asset.value))
        # no taxable income
        if abs(expenses) == 0: break

    return expenses



def unit_test():
  taxable_acct = AccountTaxable("Taxable")
  asset1 = AssetTaxable("Asset1", 1000, 0.12)
  asset2 = AssetTaxable("Asset2", 500, 0.15)
  taxable_acct.assets.append(asset1)
  taxable_acct.assets.append(asset2)

  # Test Savings
  assert(2 == len(taxable_acct.assets), lineno())
  assert(1500 == taxable_acct.value, lineno())
  taxable_acct.apply_savings(1000.0)
  assert(2500 ==  taxable_acct.value, lineno())
  assert(2 == len(taxable_acct.assets), lineno())

  # Taxable Expenses
  # basis == value in this case. no taxation zero out account.
  assert(0 == int(taxable_acct.apply_expenses(2500)), lineno())
  assert(0 == int(taxable_acct.capital_income), lineno())

  assert(0 == taxable_acct.value, lineno()) 
  assert(-2500 == taxable_acct.apply_expenses(2500), lineno())

  taxable_acct = AccountTaxable("Taxable")
  asset1 = AssetTaxable("Asset1", 1000, 0.12)
  asset2 = AssetTaxable("Asset2", 500, 0.15)
  taxable_acct.assets.append(asset1)
  taxable_acct.assets.append(asset2)
  assert(1500 == taxable_acct.value, lineno()) 

  asset1.basis = 500
  asset2.basis = 250
  assert(-1000 == taxable_acct.apply_expenses(-2500), lineno())
  assert(750 == taxable_acct.capital_income, lineno())
  assert(0 == taxable_acct.value, lineno()) 

  # Qualified Expenses
  taxable_acct = AccountTaxDeferred("Tax Deferred")
  asset1 = AssetTaxDeferred("Asset1", 1000, 0.12)
  asset2 = AssetTaxDeferred("Asset2", 500, 0.15)
  taxable_acct.assets.append(asset1)
  taxable_acct.assets.append(asset2)
  assert(1500 == taxable_acct.value, lineno()) 

  assert(-1000 == taxable_acct.apply_expenses(-2500, False), lineno())
  assert(0 == taxable_acct.capital_income, lineno())
  assert(0 == taxable_acct.value, lineno()) 

  # Tax Free
  taxable_acct = AccountTaxFree("Tax Free")
  asset1 = AssetTaxFree("Asset1", 1000, 0.12)
  asset2 = AssetTaxFree("Asset2", 500, 0.15)
  taxable_acct.assets.append(asset1)
  taxable_acct.assets.append(asset2)
  assert(1500 == taxable_acct.value, lineno()) 

  assert(-1000 == taxable_acct.apply_expenses(-2500, False), lineno())
  assert(0 == taxable_acct.capital_income, lineno())
  assert(0 == taxable_acct.value, lineno()) 

  # income test
  taxable_acct = AccountTaxFree("Tax Free")
  asset1 = AssetTaxFree("Asset1", 1000, 0.12)
  asset2 = AssetTaxFree("Asset2", 500, 0.15)
  taxable_acct.assets.append(asset1)
  taxable_acct.assets.append(asset2)
 
  taxable_acct = AccountTaxFree("Tax Free")
  asset1 = AssetTaxFree("Asset1", 1000, 0.12)
  asset2 = AssetTaxFree("Asset2", 500, 0.15)
  taxable_acct.assets.append(asset1)
  taxable_acct.assets.append(asset2)

  asset1.basis = 500
  asset2.basis = 250

  taxable_acct.apply_expenses(-2500)
  assert(taxable_acct.get_capital_income() == 750, lineno())
  assert(taxable_acct.get_capital_income() == 0, lineno())

  asset1.value = 1000
  asset1.basis = 500
  asset2.value = 500
  asset2.basis = 250

  taxable_acct = AccountTaxDeferred("whocares")
  asset1 = AssetTaxable("Asset1", 1000, 0.12)
  asset2 = AssetTaxable("Asset2", 500, 0.15)
  taxable_acct.assets.append(asset1)
  taxable_acct.assets.append(asset2)

  taxable_acct.apply_expenses(-2500, True)
  assert(0.2 * 1500 == taxable_acct.get_withdrawal_penalty(), lineno())
  assert(0 == taxable_acct.get_withdrawal_penalty(), lineno())



if __name__ == '__main__':
  unit_test()