#!/usr/bin/env python3

import enum
import logging
import re

import ansi
from type.money import Money

from person import Person
import simulation_params as sp
import simulation_results as sr
import taxman


logger = logging.getLogger(__name__)


class Verbosity(enum.IntEnum):
    SILENT = 0
    NORMAL = 1
    VERBOSE = 2


class ReportColorizer(ansi.ProgrammableColorizer):
    def __init__(self):
        super().__init__(
            [
                (re.compile('^(.+)\:$', re.MULTILINE), self.header),
                (re.compile(' *([A-Z]+[A-Z ]+)\: *'), self.title)
            ]
        )

    def header(self, match: re.match):
        return f'{ansi.fg("chateau green")}{match.group()}{ansi.reset()}'

    def title(self, match: re.match):
        s = match.group()
        if len(s) > 14:
            return f'{ansi.fg("mint green")}{ansi.bold()}{ansi.underline()}{s}{ansi.reset()}'
        return s


class Simulation(object):
    def __init__(
            self,
            params: sp.SimulationParameters,
    ):
        self.year = 0
        self.scott_age = 0
        self.lynn_age = 0
        self.alex_age = 0
        self.accounts = [x for x in params.initial_account_states]
        self.taxes = taxman.TaxCollector(
            params.federal_ordinary_income_tax_brackets,
            params.federal_dividends_and_long_term_gains_brackets
        )
        self.params = params
        self.returns_and_expenses = params.returns_and_expenses

    def do_rmd_withdrawals(self) -> Money:
        """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() == Person.SCOTT:
                rmd = x.do_rmd_withdrawal(self.scott_age, self.taxes)
                total_withdrawn += rmd
            elif x.has_rmd() and (x.get_owner() == Person.LYNN or x.get_owner() == Person.SHARED):
                rmd = x.do_rmd_withdrawal(self.lynn_age, self.taxes)
                total_withdrawn += rmd
        return total_withdrawn

    def go_find_money(self, amount: Money):
        """Look through accounts and try to find amount using some heuristics
           about where to withdraw money first."""
        logger.debug(f'Simulation: looking for {amount}...')

        # Try brokerage accounts first
        for x in self.accounts:
            if not x.is_age_restricted():
                if x.get_balance() >= amount:
                    logger.debug(
                        f'Simulation: withdrawing {amount} from {x.get_name()}'
                    )
                    x.withdraw(amount, self.taxes)
                    return

                elif x.get_balance() > 0.01 and x.get_balance() < amount:
                    money_available = x.get_balance()
                    logger.debug(
                        f'Simulation: withdrawing {money_available} from {x.get_name()} ' +
                        '(and exhausting the account)'
                    )
                    x.withdraw(money_available, self.taxes)
                    amount -= money_available

        # 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
            ):
                if x.get_balance() >= amount:
                    logging.debug(
                        f'Simulation: withdrawing {amount} from {x.get_name()}'
                    )
                    x.withdraw(amount, self.taxes)
                    return

                elif x.get_balance() > 0.01 and x.get_balance() < amount:
                    money_available = x.get_balance()
                    logger.debug(
                        f'Simulation: withdrawing {money_available} from {x.get_name()} ' +
                        '(and exhausting the account)'
                    )
                    x.withdraw(money_available, self.taxes)
                    amount -= money_available

            if (
                    x.is_age_restricted() and
                    x.has_rmd() and
                    x.belongs_to_scott() and
                    self.scott_age > 60
            ):
                if x.get_balance() >= amount:
                    logger.debug(
                        f'Simulation: withdrawing {amount} from {x.get_name()}'
                    )
                    x.withdraw(amount, self.taxes)
                    return

                elif x.get_balance() > 0.01 and x.get_balance() < amount:
                    money_available = x.get_balance()
                    logger.debug(
                        f'Simulation: withdrawing {money_available} from {x.get_name()} ' +
                        '(and exhausting the account)'
                    )
                    x.withdraw(money_available, self.taxes)
                    amount -= money_available

        # 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
            ):
                if x.get_balance() >= amount:
                    logger.debug(
                        f'Simulation: withdrawing {amount} from {x.get_name()}'
                    )
                    x.withdraw(amount, self.taxes)
                    return

                elif x.get_balance() > 0.01 and x.get_balance() < amount:
                    money_available = x.get_balance()
                    logger.debug(
                        f'Simulation: withdrawing {money_available} from {x.get_name()} ' +
                        '(and exhausting the account)'
                    )
                    x.withdraw(money_available, self.taxes)
                    amount -= money_available

            if (
                    x.is_age_restricted() and
                    x.has_roth() and
                    x.belongs_to_scott() and
                    self.scott_age > 60
            ):
                if x.get_balance() >= amount:
                    logger.debug(
                        f'Simulation: withdrawing {amount} from {x.get_name()}'
                    )
                    x.withdraw(amount, self.taxes)
                    return

                elif x.get_balance() > 0.01 and x.get_balance() < amount:
                    money_available = x.get_balance()
                    logger.debug(
                        f'Simulation: withdrawing {money_available} from {x.get_name()} ' +
                        '(and exhausting the account)'
                    )
                    x.withdraw(money_available, self.taxes)
                    amount -= money_available
        raise Exception("Unable to find enough money this year!")

    def get_social_security(self) -> Money:
        total_benefit = Money(0)
        if self.scott_age >= self.params.scott_social_security_age:
            total_benefit += self.params.scott_annual_social_security_dollars
            self.taxes.record_ordinary_income(
                self.params.scott_annual_social_security_dollars
            )
        if self.lynn_age >= self.params.lynn_social_security_age:
            total_benefit += self.params.lynn_annual_social_security_dollars
            self.taxes.record_ordinary_income(
                self.params.lynn_annual_social_security_dollars
            )
        return total_benefit

    def do_opportunistic_roth_conversions(self) -> Money:
        desired_conversion_amount = (
            self.taxes.how_many_more_dollars_can_we_earn_without_changing_tax_rate()
        )
        money_converted = Money(0)
        for x in self.accounts:
            if x.has_roth():
                amount = x.do_roth_conversion(desired_conversion_amount)
                money_converted += amount
                desired_conversion_amount -= amount
        return money_converted

    def simulate(
            self,
            silent: Verbosity = Verbosity.NORMAL
    ) -> sr.SimulationResults:
        if silent is not Verbosity.SILENT:
            with ReportColorizer():
                print(self.params)

        results = sr.SimulationResults.create()
        try:
            for year in range(2022, 2069):
                self.year = year
                self.scott_age = year - 1974
                self.lynn_age = year - 1964
                self.alex_age = year - 2005

                # Computed money needed this year based on inflated budget.
                money_needed = self.params.annual_expenses

                # When Alex is in college, we need $60K more per year.
                if self.alex_age > 18 and self.alex_age <= 22:
                    money_needed += Money(60000)

                if silent is Verbosity.VERBOSE:
                    print(
                        f'\nYear: {year}, estimated annual expenses {money_needed} -----------------\n'
                    )
                    print(
                        f'Scott is {self.scott_age}, Lynn is {self.lynn_age} and Alex is {self.alex_age}'
                    )

                # Print out how much money is in each account +
                # overall net worth.
                total = 0
                if silent is not Verbosity.VERBOSE:
                    for x in self.accounts:
                        total += x.get_balance()
                else:
                    for x in self.accounts:
                        total += x.get_balance()
                        print('  %-50s: %18s' % (x.get_name(), x.get_balance()))
                    print("  %-50s: %18s" % ("TOTAL NET WORTH", total))

                # Max
                if total > results.max_networth:
                    results.max_networth = total
                    results.max_networth_year = year

                # 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()
                if rmds > 0:
                    logger.debug(
                        f'Simulation: satisfied {rmds} of RMDs from age-restricted accounts.'
                    )
                    total_income += rmds
                    money_needed -= rmds
                    if rmds > results.max_rmd:
                        results.max_rmd = rmds
                        results.max_rmd_year = year

                # When we reach a certain age we are eligible for SS
                # payments.
                ss = self.get_social_security()
                results.total_social_security += ss
                if ss > 0:
                    logger.debug(
                        f'Simulation: received {ss} from social security.'
                    )
                    total_income += ss
                    money_needed -= ss

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

                # Maybe do some opportunistic Roth conversions.
                taxes_due = self.taxes.approximate_annual_taxes()
                old_taxes_due = taxes_due
                tax_rate = self.taxes.approximate_annual_tax_rate()
                if tax_rate <= 0.16 and year <= 2035:
                    rollover = self.do_opportunistic_roth_conversions()
                    if rollover > results.max_rollover:
                        results.max_rollover = rollover
                        results.max_rollover_year = year
                    results.total_rollover += rollover

                    # These conversions will affect taxes due.
                    # Recompute them now.
                    if rollover > 0:
                        taxes_due = self.taxes.approximate_annual_taxes()
                        tax_rate = self.taxes.approximate_annual_tax_rate()
                        logger.debug(
                            f'Simulation: rolled {rollover} from pretax into Roth which increased '
                            f'taxes by {taxes_due - old_taxes_due}'
                        )

                # 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 > results.max_taxes:
                    results.max_taxes = taxes_due
                    results.max_taxes_year = year
                if taxes_due > 0:
                    logger.debug(
                        f'Simulation: estimated {taxes_due} federal tax liability (tax rate={tax_rate})'
                    )
                    self.go_find_money(taxes_due)
                else:
                    logger.debug('Simulation: estimating no tax burden this year')
                self.taxes.record_taxes_paid_and_reset_for_next_year(year)

                # Inflation and appreciation.  Note also: SS benefits
                # currently are increased annually to adjust for
                # inflation.  Do that here too.
                inflation = self.returns_and_expenses.get_annual_inflation_rate(year)
                roi = self.returns_and_expenses.get_annual_investment_return_rate(year)
                ss = self.returns_and_expenses.get_annual_social_security_increase_rate(year)
                self.params.annual_expenses *= inflation
                total = 0
                for x in self.accounts:
                    total += x.appreciate(roi)
                if self.scott_age >= self.params.scott_social_security_age:
                    self.params.scott_annual_social_security_dollars *= ss
                if self.lynn_age >= self.params.lynn_social_security_age:
                    self.params.lynn_annual_social_security_dollars *= ss
                msg = (
                    f'In {year}, '
                    f'inflation={inflation.__repr__(relative=True)}, ' +
                    f'ROI={roi.__repr__(relative=True)}, ' +
                    f'net worth={total}'
                )
                logger.debug(msg)
                if silent is not Verbosity.SILENT:
                    print(msg)

        except Exception as e:
            if silent is not Verbosity.SILENT:
                logger.exception(e)
                print(f'**** Exception, out of money in {year}? ****')
            results.success = False
        else:
            results.success = True

        # Final report
        results.end_year = self.year + 1
        results.final_networth = Money(0)
        if results.success:
            for x in self.accounts:
                results.final_networth += x.get_balance()
        results.final_taxes = self.taxes
        results.final_account_state = self.accounts
        if silent is not Verbosity.SILENT:
            with ReportColorizer():
                print(results)
        return results