import constants
import utils

class account(object):
    """This defines an account base class which inherits from object, the
       default base class in python.  This is not a 100% abstract class, it
       includes some default fields (like balance and name) and methods that
       define behaviors (withdraw, appreciate, etc...).  Several subclasses
       are defined below."""

    def __init__(self, name, owner):
        """Constructor for account object takes a name and owner"""
        self.name = name
        self.owner = owner

    def get_name(self):
        return self.name

    def get_owner(self):
        return self.owner

    def belongs_to_scott(self):
        return self.get_owner() == constants.SCOTT

    def belongs_to_lynn(self):
        return self.get_owner() == constants.LYNN

    def get_balance(self):
        pass

    def appreciate(self, multiplier):
        """Grow the balance using rate."""
        pass

    def withdraw(self, amount, taxes):
        """Withdraw money from the account if possible.  Throws otherwise.
           Money should be registered with taxes so that we can approximate
           taxes due, too."""
        pass

    def deposit(self, amount):
        pass

    def is_age_restricted(self):
        """Is this account age restricted.  Subclasses should implement."""
        pass

    def has_rmd(self):
        """Does this account have a required minimum distribution?  Sub-
           classes should implement."""
        pass

    def do_rmd_withdrawal(self, owner_age, taxes):
        pass

    def has_roth(self):
        pass

    def dump_final_report(self):
        print "  " + self.name + ":"

class age_restricted_tax_deferred_account(account):
    """Account object used to represent an age-restricted tax deferred
       account such as a 401(k), IRA or annuity."""

    def __init__(self, total_balance, roth_amount_subset, name, owner):
        """C'tor for this class takes a roth_amount, the number of dollars
           in the account that are in-plan Roth and can be taken out tax
           free.  It keeps this pile separate from the pretax pile and
           tries to estimate taxes."""
        self.roth = roth_amount_subset
        self.pretax = total_balance - roth_amount_subset
        self.total_roth_withdrawals = 0
        self.total_pretax_withdrawals = 0
        self.total_investment_gains = 0
        self.total_roth_conversions = 0
        self.initial_balance = self.get_balance()

        # This calls the super class' c'tor.
        account.__init__(self, name, owner)

    # These accounts are all age-restricted
    def is_age_restricted(self):
        return True

    def withdraw(self, amount, taxes):
        """Assume that money withdrawn from this account will be a mixture
           of pretax funds (which count as ordinary income) and Roth funds
           (which are available tax-free).  Assume that the ratio of pretax
           to Roth in this overall account determines the amount from each
           partition in this withdrawal."""
        balance = self.get_balance()
        if balance < amount:
            raise Exception("Insufficient funds")

        ratio = float(self.roth) / float(balance)
        roth_part = amount * ratio
        pretax_part = amount - roth_part
        if roth_part > 0:
            self.roth -= roth_part
            self.total_roth_withdrawals += roth_part
            print "## Satisfying %s from %s with Roth money." % (utils.format_money(roth_part), self.name)
            taxes.record_roth_income(roth_part)

        self.pretax -= pretax_part
        self.total_pretax_withdrawals += pretax_part
        print "## Satisfying %s from %s with pre-tax money." % (utils.format_money(pretax_part), self.name)
        taxes.record_ordinary_income(pretax_part)

    def deposit(self, amount):
        self.pretax_part += amount

    def appreciate(self, multiplier):
        """In this class we basically ignore the balance field in favor of
           just using pretax and roth so that we can track them separately."""
        old_pretax = self.pretax
        self.pretax *= multiplier
        self.total_investment_gains += self.pretax - old_pretax
        old_roth = self.roth
        self.roth *= multiplier
        self.total_investment_gains += self.roth - old_roth

    def get_balance(self):
        """In this class we keep the balance in two parts."""
        return self.pretax + self.roth

    def has_rmd(self):
        return True

    # Used to compute RMDs.  Source:
    # https://www.irs.gov/publications/p590b#en_US_2018_publink1000231258
    def get_actuary_number_years_to_live(self, age):
        if age < 70:
            return 27.4 + age
        elif age == 70:
            return 27.4
        elif age == 71:
            return 26.5
        elif age == 72:
            return 25.6
        elif age == 73:
            return 24.7
        elif age == 74:
            return 23.8
        elif age == 75:
            return 22.9
        elif age == 76:
            return 22.0
        elif age == 77:
            return 21.2
        elif age == 78:
            return 20.3
        elif age == 79:
            return 19.5
        elif age == 80:
            return 18.7
        elif age == 81:
            return 17.9
        elif age == 82:
            return 17.1
        elif age == 83:
            return 16.3
        elif age == 84:
            return 15.5
        elif age == 85:
            return 14.8
        elif age == 86:
            return 14.1
        elif age == 87:
            return 13.4
        elif age == 88:
            return 12.7
        elif age == 89:
            return 12.0
        elif age == 90:
            return 11.4
        elif age == 91:
            return 10.8
        elif age == 92:
            return 10.2
        elif age == 93:
            return 9.6
        elif age == 94:
            return 9.1
        elif age == 95:
            return 8.6
        elif age == 96:
            return 8.1
        elif age == 97:
            return 7.6
        elif age == 98:
            return 7.1
        elif age == 99:
            return 6.7
        elif age == 100:
            return 6.3
        elif age == 101:
            return 5.9
        elif age == 102:
            return 5.5
        elif age == 103:
            return 5.2
        elif age == 104:
            return 4.9
        elif age == 105:
            return 4.5
        elif age == 106:
            return 4.2
        elif age == 107:
            return 3.9
        elif age == 108:
            return 3.7
        elif age == 109:
            return 3.4
        elif age == 110:
            return 3.1
        elif age == 111:
            return 2.9
        elif age == 112:
            return 2.6
        elif age == 113:
            return 2.4
        elif age == 114:
            return 2.1
        else:
            return 1.9

    def do_rmd_withdrawal(self, owner_age, taxes):
        balance = self.get_balance()
        if balance > 0 and owner_age >= 72:
            rmd_factor = self.get_actuary_number_years_to_live(owner_age)
            amount = balance / rmd_factor
            self.withdraw(amount, taxes)
            return amount
        return 0

    def has_roth(self):
        return True

    def do_roth_conversion(self, amount):
        if amount <= 0:
            return 0
        if self.pretax >= amount:
            self.roth += amount
            self.pretax -= amount
            self.total_roth_conversions += amount
            print "## Executed pre-tax --> Roth conversion of %s in %s" % (
                utils.format_money(amount),
                self.get_name())
            return amount
        elif self.pretax > 0:
            actual_amount = self.pretax
            self.roth += actual_amount
            self.pretax = 0
            self.total_roth_conversions += actual_amount
            print "## Executed pre-tax --> Roth conversion of %s in %s" % (
                utils.format_money(actual_amount),
                self.get_name())
            return actual_amount
        return 0

    def dump_final_report(self):
        super(age_restricted_tax_deferred_account, self).dump_final_report()
        print "    {:<50}: {:>14}".format("Initial balance",
                                          utils.format_money(self.initial_balance))
        print "    {:<50}: {:>14}".format("Ending balance",
                                        utils.format_money(self.get_balance()))
        print "    {:<50}: {:>14}".format("Total investment gains",
                                        utils.format_money(self.total_investment_gains))
        print "    {:<50}: {:>14}".format("Total Roth withdrawals",
                                        utils.format_money(self.total_roth_withdrawals))
        print "    {:<50}: {:>14}".format("Total pre-tax withdrawals",
                                        utils.format_money(self.total_pretax_withdrawals))
        print "    {:<50}: {:>14}".format("Total pre-tax converted to Roth",
                                        utils.format_money(self.total_roth_conversions))

class age_restricted_roth_account(age_restricted_tax_deferred_account):
    """This is an object to represent a Roth account like a Roth IRA.  All
       money in here is tax free.  Most of the base account class works here
       including the implementation of withdraw() which says that none of the
       money withdrawn was taxable."""

    def __init__(self, total_balance, name, owner):
        """C'tor for this class takes a roth_amount, the number of dollars
           in the account that are in-plan Roth and can be taken out tax
           free.  It keeps this pile separate from the pretax pile and
           tries to estimate taxes."""
        age_restricted_tax_deferred_account.__init__(
            self, total_balance, 0, name, owner)

    # Override
    def has_rmd(self):
        return False

    # Override
    def do_rmd_withdrawal(self, owner_age, taxes):
        raise Exception("This account has no RMDs")

    # Override
    def do_roth_conversion(self, amount):
        return 0

class brokerage_account(account):
    """A class to represent money in a taxable brokerage account."""

    def __init__(self, total_balance, cost_basis, name, owner):
        """The c'tor of this class partitions balance into three pieces:
           the cost_basis (i.e. how much money was invested in the account),
           the short_term_gain (i.e. appreciation that has been around for
           less than a year) and long_term_gain (i.e. appreciation that has
           been around for more than a year).  We separate these because
           taxes on long_term_gain (and qualified dividends, which are not
           modeled) are usually lower than short_term_gain.  Today those
           are taxed at 15% and as ordinary income, respectively."""
        self.cost_basis = float(cost_basis)
        self.short_term_gain = 0.0
        self.long_term_gain = float(total_balance - cost_basis)
        self.total_cost_basis_withdrawals = 0
        self.total_long_term_gain_withdrawals = 0
        self.total_short_term_gain_withdrawals = 0
        self.total_investment_gains = 0
        self.initial_balance = self.get_balance()

        # Call the base class' c'tor as well to set up the name.
        account.__init__(self, name, owner)

    def withdraw(self, amount, taxes):
        """Override the base class' withdraw implementation since we're
           dealing with three piles of money instead of one.  When you sell
           securities to get money out of this account the gains are taxed
           (and the cost_basis part isn't).  Assume that the ratio of
           cost_basis to overall balance can be used to determine how much
           of the withdrawal will be taxed (and how)."""
        balance = self.get_balance()
        if balance < amount:
            raise Exception("Insufficient funds")

        ratio = self.cost_basis / balance
        invested_capital_part = amount * ratio
        gains_part = amount - invested_capital_part

        if self.cost_basis >= invested_capital_part:
            self.cost_basis -= invested_capital_part
            self.total_cost_basis_withdrawals += invested_capital_part
            print "## Satisfying %s from %s as cost basis." % (utils.format_money(invested_capital_part), self.name)
            self.withdraw_from_gains(gains_part, taxes)
        else:
            self.withdraw_from_gains(amount, taxes)

    def withdraw_from_gains(self, amount, taxes):
        """Withdraw some money from gains.  Prefer the long term ones if
           possible."""
        if amount > (self.long_term_gain + self.short_term_gain):
            raise Exception("Insufficient funds")

        if self.long_term_gain >= amount:
            self.long_term_gain -= amount
            self.total_long_term_gain_withdrawals += amount
            print "## Satisfying %s from %s as long term gains." % (
                utils.format_money(amount),
                self.name)
            taxes.record_dividend_or_long_term_gain(amount)
            return

        else:
            print "## Satisfying %s from %s as long term gains (exhausting all current long term gains in account)." % (
                utils.format_money(self.long_term_gain),
                self.name)
            amount -= self.long_term_gain
            self.total_long_term_gain_withdrawals += self.long_term_gain
            taxes.record_dividend_or_long_term_gain(self.long_term_gain)
            self.long_term_gain = 0
            # Get the rest of amount from short term gains...

        self.short_term_gain -= amount
        self.total_short_term_gain_withdrawals += amount
        print "## Satisfying %s from %s as short term gains." % (
            utils.format_money(amount),
            self.name)
        taxes.record_short_term_gain(amount)

    def deposit(self, amount):
        self.cost_basis += amount

    def get_balance(self):
        """We're ignoring the base class' balance field in favor of tracking
           it as three separate partitions of money."""
        return self.cost_basis + self.long_term_gain + self.short_term_gain

    def appreciate(self, multiplier):
        """Appreciate... another year has passed so short_term_gains turn into
           long_term_gains and the appreciation is our new short_term_gains."""
        balance = self.get_balance()
        balance *= multiplier
        gain_or_loss = balance - self.get_balance()
        self.total_investment_gains += gain_or_loss
        self.long_term_gain += self.short_term_gain
        self.short_term_gain = gain_or_loss

    def is_age_restricted(self):
        return False

    def has_rmd(self):
        return False

    def has_roth(self):
        return False

    def do_roth_conversion(self, amount):
        return 0

    def dump_final_report(self):
        super(brokerage_account, self).dump_final_report()
        print "    {:<50}: {:>14}".format("Initial balance",
                                          utils.format_money(self.initial_balance))
        print "    {:<50}: {:>14}".format("Ending balance",
                                        utils.format_money(self.get_balance()))
        print "    {:<50}: {:>14}".format("Total investment gains",
                                        utils.format_money(self.total_investment_gains))
        print "    {:<50}: {:>14}".format("Total cost basis withdrawals",
                                        utils.format_money(self.total_cost_basis_withdrawals))
        print "    {:<50}: {:>14}".format("Total long term gain withdrawals",
                                        utils.format_money(self.total_long_term_gain_withdrawals))
        print "    {:<50}: {:>14}".format("Total short term gain withdrawals",
                                        utils.format_money(self.total_short_term_gain_withdrawals))