#!/usr/local/bin/python

import sys

from accounts import *
import constants
from parameters import parameters
from tax_brackets import tax_brackets
from tax_collector import tax_collector
import utils
import secrets

class simulation(object):
    def __init__(self, parameters, accounts):
        """Initialize simulation parameters and starting account balances"""
        self.params = parameters
        self.accounts = accounts
        self.lynn_age = 0
        self.scott_age = 0
        self.alex_age = 0
        self.year = 0

    def do_rmd_withdrawals(self, taxes):
        """Determine if any account that has RMDs will require someone to
           take money out of it this year and, if so, how much.  Then do
           the withdrawal."""
        total_withdrawn = 0
        for x in self.accounts:
            if x.has_rmd() and x.get_owner() == constants.SCOTT:
                rmd = x.do_rmd_withdrawal(self.scott_age, taxes)
                total_withdrawn += rmd
            elif x.has_rmd() and x.get_owner() == constants.LYNN:
                rmd = x.do_rmd_withdrawal(self.lynn_age, taxes)
                total_withdrawn += rmd
        return total_withdrawn

    def go_find_money(self, amount_needed, taxes):
        """Look through accounts and try to find amount using some heuristics
           about where to withdraw money first."""

        # Try brokerage accounts first
        for x in self.accounts:
            if not x.is_age_restricted():
                amount_to_withdraw = min(amount_needed, x.get_balance())
                print "## Withdrawing %s from %s" % (utils.format_money(amount_to_withdraw), x.get_name())
                x.withdraw(amount_to_withdraw, taxes)
                amount_needed -= amount_to_withdraw
                if amount_needed <= 0: return

        # Next try age restircted accounts
        for x in self.accounts:
            if (x.is_age_restricted() and
                x.has_rmd() and
                ((x.belongs_to_lynn() and self.lynn_age > 60) or
                 (x.belongs_to_scott() and self.scott_age > 60))):

                amount_to_withdraw = min(amount_needed, x.get_balance())
                print "## Withdrawing %s from %s" % (utils.format_money(amount_to_withdraw), x.get_name())
                x.withdraw(amount_to_withdraw, taxes)
                amount_needed -= amount_to_withdraw
                if amount_needed <= 0: return

        # Last try Roth accounts
        for x in self.accounts:
            if (x.is_age_restricted() and
                x.has_roth() and
                ((x.belongs_to_lynn() and self.lynn_age > 60) or
                 (x.belongs_to_scott() and self.scott_age > 60))):

                amount_to_withdraw = min(amount_needed, x.get_balance())
                print "## Withdrawing %s from %s" % (utils.format_money(amount_to_withdraw), x.get_name())
                x.withdraw(amount_to_withdraw, taxes)
                amount_needed -= amount_to_withdraw
                if amount_needed <= 0: return
        raise Exception("Unable to find enough money this year, still need %s more!" % utils.format_money(amount_needed))

    def get_social_security(self,
                            scott_annual_social_security_dollars,
                            lynn_annual_social_security_dollars,
                            taxes):
        """Figure out if Scott and/or Lynn are taking social security at
           their present age in the simulation and, if so, how much their
           annual benefit should be."""
        total_benefit = 0
        if self.scott_age >= self.params.get_initial_social_security_age(constants.SCOTT):
            total_benefit += scott_annual_social_security_dollars
            taxes.record_ordinary_income(scott_annual_social_security_dollars)
        if self.lynn_age >= self.params.get_initial_social_security_age(constants.LYNN):
            total_benefit += lynn_annual_social_security_dollars
            taxes.record_ordinary_income(lynn_annual_social_security_dollars)
        return total_benefit

    def do_opportunistic_roth_conversions(self, taxes):
        """Roll over money from pretax 401(k)s or IRAs into Roth."""
        desired_conversion_amount = (
            taxes.how_many_more_dollars_can_we_earn_without_changing_tax_rate(
                self.params.get_federal_ordinary_income_tax_brackets()))
        if desired_conversion_amount > 0:
            money_converted = 0
            for x in self.accounts:
                if x.has_roth():
                    amount = x.do_roth_conversion(desired_conversion_amount, taxes)
                    money_converted += amount
                    desired_conversion_amount -= amount
            if money_converted > 0:
                return True
        return False

    def dump_report_header(self):
        self.params.dump()

    def dump_annual_header(self, money_needed):
        print "\nYear: %d, estimated annual expenses %s ---------------\n" % (
            self.year,
            utils.format_money(money_needed))
        print "Scott is %d, Lynn is %d and Alex is %d.\n" % (
            self.scott_age, self.lynn_age, self.alex_age)

        # Print out how much money is in each account + overall net worth.
        total = 0
        for x in self.accounts:
            total += x.get_balance()
            print "{:<50}: {:>14}".format(x.get_name(),
                                          utils.format_money(x.get_balance()))
        print "{:<50}: {:>14}\n".format("TOTAL", utils.format_money(total))

    def dump_final_report(self, taxes):
        print "\nAGGREGATE STATS FINAL REPORT:"
        for x in self.accounts:
            x.dump_final_report()
        taxes.dump_final_report()

    def run(self):
        """Run the simulation!"""
        self.dump_report_header()

        taxes = tax_collector()
        adjusted_annual_expenses = self.params.get_initial_annual_expenses()
        adjusted_scott_annual_social_security_dollars = (
            self.params.get_initial_social_security_benefit(constants.SCOTT))
        adjusted_lynn_annual_social_security_dollars = (
            self.params.get_initial_social_security_benefit(constants.LYNN))

        try:
            for self.year in xrange(2020, 2080):
                self.scott_age = self.year - 1974
                self.lynn_age = self.year - 1964
                self.alex_age = self.year - 2005

                # Computed money needed this year based on inflated budget.
                money_needed = adjusted_annual_expenses

                # When Alex is in college, we need $50K more per year.
                if self.alex_age > 18 and self.alex_age <= 22:
                    money_needed += 50000

                self.dump_annual_header(money_needed)

                # Now, figure out how to find money to pay for this year
                # and how much of it is taxable.
                total_income = 0

                # When we reach a certain age we have to take RMDs from
                # some accounts.  Handle that here.
                rmds = self.do_rmd_withdrawals(taxes)
                if rmds > 0:
                    print "## Satisfied %s of RMDs from age-restricted accounts." % utils.format_money(rmds)
                    total_income += rmds
                    money_needed -= rmds

                # When we reach a certain age we are eligible for SS
                # payments.
                ss = self.get_social_security(adjusted_scott_annual_social_security_dollars,
                                              adjusted_lynn_annual_social_security_dollars,
                                              taxes)
                if ss > 0:
                    print "## Social security paid %s" % utils.format_money(ss)
                    total_income += ss
                    money_needed -= ss

                # If we still need money, try to go find it.
                if money_needed > 0:
                    self.go_find_money(money_needed, taxes)
                    total_income += money_needed
                    money_needed = 0

                look_for_conversions = True
                while True:
                    # Maybe do some opportunistic Roth conversions.
                    taxes_due = taxes.approximate_taxes(
                        self.params.get_federal_standard_deduction(),
                        self.params.get_federal_ordinary_income_tax_brackets(),
                        self.params.get_federal_dividends_and_long_term_gains_income_tax_brackets())
                    total_income = taxes.get_total_income()
                    tax_rate = float(taxes_due) / float(total_income)
                    print "INCOME: %s, TAXES: %s\n" % (utils.format_money(total_income), utils.format_money(taxes_due))

                    if (look_for_conversions and
                        tax_rate <= 0.14 and
                        taxes_due < 20000 and
                        self.year <= 2035):

                        look_for_conversions = self.do_opportunistic_roth_conversions(taxes)
                        # because these conversions affect taxes, spin
                        # once more in the loop and recompute taxes
                        # due and tax rate.
                    else:
                        break

                # Pay taxes_due by going to find more money.  This is a
                # bit hacky since withdrawing more money to cover taxes
                # will, in turn, cause taxable income.  But I think taxes
                # are low enough and this simulation is rough enough that
                # we can ignore this.
                if taxes_due > 0:
                    print "## Estimated federal tax due: %s (this year's tax rate=%s)" % (
                        utils.format_money(taxes_due),
                        utils.format_rate(tax_rate))
                    self.go_find_money(taxes_due, taxes)
                else:
                    print "## No federal taxes due this year!"
                taxes.record_taxes_paid_and_reset_for_next_year(taxes_due)

                # Inflation and appreciation:
                #   * Cost of living increases
                #   * Social security benefits increase
                #   * Tax brackets are adjusted for inflation
                inflation_multiplier = self.params.get_average_inflation_multiplier()
                adjusted_annual_expenses *= inflation_multiplier
                for x in self.accounts:
                    x.appreciate(self.params.get_average_investment_return_multiplier())
                if self.scott_age >= self.params.get_initial_social_security_age(constants.SCOTT):
                    adjusted_scott_annual_social_security_dollars *= self.params.get_average_social_security_multiplier()
                if self.lynn_age >= self.params.get_initial_social_security_age(constants.LYNN):
                    adjusted_lynn_annual_social_security_dollars *= self.params.get_average_social_security_multiplier()

                self.params.get_federal_ordinary_income_tax_brackets().adjust_with_multiplier(inflation_multiplier)
                self.params.get_federal_dividends_and_long_term_gains_income_tax_brackets().adjust_with_multiplier(inflation_multiplier)
        except:
            print "Ran out of money!!!"
            pass

        finally:
            self.dump_final_report(taxes)

# main
params = parameters().with_default_values()
accounts = secrets.accounts
s = simulation(params, accounts)
s.run()