import pandas as pd
import numpy as np
import datetime as dt
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
plt.style.use('fivethirtyeight')

def calc_payment_breakdown(period_rate, period, n_periods, principal):
    ipmt = np.ipmt(period_rate, period, n_periods, principal)
    ppmt = np.ppmt(period_rate, period, n_periods, principal)
    return ipmt, ppmt

def gen_schedule(loan_amount, rate, start_date, length, pn=12):
    n_periods = length
    period_rate = rate / pn

    tpmt = np.pmt(period_rate, n_periods, loan_amount)
    period_values = [(i, tpmt, *calc_payment_breakdown(period_rate, i, n_periods, loan_amount))
                     for i in range(1, n_periods+1)]
    ix = pd.date_range(start_date, periods=n_periods, freq='MS')

    cols = ['n', 'Payment', 'Interest', 'Principal']
    df = pd.DataFrame(period_values, columns=cols, index=ix)
    df['Balance'] = (loan_amount - (-1*df.Principal.cumsum()))
    df['Equity'] = df['Principal'].abs().cumsum()

    return df

def calc_closing_costs(rate, credits, loan_amount=359572):
    loan_costs = 2122
    # other costs
    taxes_fees = 120
    # prepaids
    prepaid_prop_tax = 4801
    loan_amnt_15days = loan_amount / 360 * 15
    # escrow
    homeowners_insurance = 291
    property_taxes = 3201

    rate_indepedent = (loan_costs +
                       taxes_fees +
                       prepaid_prop_tax +
                       homeowners_insurance +
                       property_taxes)
    return rate_indepedent + loan_amnt_15days*rate + credits

def calc_comparison(current, new, closing=0):
    merged = pd.merge(current, new, how='outer', left_index=True, right_index=True)

    cash_diff = (merged.Payment_x.abs().fillna(0).cumsum() -
                 (merged.Payment_y.abs().fillna(0).cumsum() + closing))
    principal_diff = (merged.Principal_x.abs().fillna(0).cumsum() -
                      merged.Principal_y.abs().fillna(0).cumsum())
    total_savings = cash_diff - principal_diff

    return cash_diff, principal_diff, total_savings

def calc_payments(current, new, term):
    current_pmi = 182.16
    total_escrow = 1125.86
    escrow = total_escrow - current_pmi
    new_pmi = {20: 54, 30: 93}

    current_payment = current.Payment.abs().mean() + total_escrow
    new_payment = new.Payment.abs().mean() + escrow + new_pmi[term]

    return current_payment, new_payment

def plot_savings(term, rate, savings, axes):
    ax1, ax2 = axes
    label = f'{term} yr @ {rate}%'
    ax1.scatter(savings.index, savings, marker='.', label=label)
    ax1.legend(fontsize=12, markerscale=2)

    m = (savings.index ==  start_date) | ((savings.index.year <= 2030) & (savings.index.month == 4))
    ax2.scatter(savings.index[m], savings[m], label=label)
    ax2.legend(fontsize=12)

    for ax in (ax1, ax2):
        ylbls = [f'${y/1000:0.0f}K' for y in ax.get_yticks()]
        ax.set_yticklabels(ylbls)
        ax.set_ylabel('Total Savings', fontsize=18)
        ax.set_xlabel('Date', fontsize=18)

## CONSTANTS
base_schedule = gen_schedule(370500, 0.04625, '2018-05-01', 360)
past_schedule = base_schedule.loc[:'2020-04-01']
current_schedule = base_schedule.loc['2020-05-01':]

current_amount = past_schedule.Balance.values[-2]

# rate options
options = [(3.125, 2701, 30),
           (3.25, 162, 30),
           (3.375, -1762, 30),
           (3.125, 1571, 20),
           (3.25, -762, 20),
           (3.375, -2636, 20),
          ]
escrow_balance = 5788
start_date = '2020-05-01'

##########################################
## Closing Paid in Cash at Close
##########################################

f, (ax1, ax2) = plt.subplots(2, figsize=(12, 18))
for (rate, points, term) in options:
    rate /= 100

    closing = calc_closing_costs(rate, points)

    new_schedule = gen_schedule(current_amount, rate, start_date, term*12)
    cash, equity, savings = calc_comparison(current_schedule, new_schedule, closing)
    current_payment, new_payment = calc_payments(current_schedule, new_schedule, term)
    payment_diff = current_payment - new_payment

    breakeven = savings.loc[savings > 0].head(1)

    rate *= 100
    print(f'Rate {rate}%, Points ${points}, Term {term}')
    print(f'\tClosing Due: ${closing:0.2f}, Less Escrow Balance: ${closing-escrow_balance:0.2f}')
    print(f'\tRecoup Closing On: {breakeven.index[0].strftime("%Y-%m-%d")}, ${breakeven.values[0]:0.2f}')
    print(f'\tMonthly Payments: Current ${current_payment:0.2f}, New ${new_payment:0.2f}, '
          f'Difference:${payment_diff:0.2f}')
    print()
    plot_savings(term, rate, savings, (ax1, ax2))
    f.suptitle('Refinance Rate Comparison\nClosing Costs Paid in Full at Close', fontsize=18, y=0.92)

##########################################
## Closing Built in to Loan
##########################################

f, (ax1, ax2) = plt.subplots(2, figsize=(12, 18))
for (rate, points, term) in options:
    rate /= 100

    # cap addition to loan at $10K
    closing = calc_closing_costs(rate, points)
    if closing > 10000:
        to_add = 10000
    else:
        to_add = closing
    closing -= to_add

    new_schedule = gen_schedule(current_amount + to_add, rate, start_date, term*12)
    cash, equity, savings = calc_comparison(current_schedule, new_schedule, closing)
    current_payment, new_payment = calc_payments(current_schedule, new_schedule, term)
    payment_diff = current_payment - new_payment

    breakeven = savings.loc[savings > 0].head(1)

    rate *= 100
    print(f'Rate {rate}%, Points ${points}, Term {term}')
    print(f'\tClosing Due: ${closing:0.2f}, Less Escrow Balance: ${closing-escrow_balance:0.2f}')
    print(f'\tRecoup Closing On: {breakeven.index[0].strftime("%Y-%m-%d")}, ${breakeven.values[0]:0.2f}')
    print(f'\tMonthly Payments: Current ${current_payment:0.2f}, New ${new_payment:0.2f}, '
          f'Difference:${payment_diff:0.2f}')
    print()
    plot_savings(term, rate, savings, (ax1, ax2))
    f.suptitle('Refinance Rate Comparison\nClosing Costs Added to Loan ($10K max)', fontsize=18, y=0.92)