from portfolio import np, pretty_num, portfolio, INFLATION_RATE, set_INIT_AGE

import gradio as gr

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure

def net_worth_calculator(Annual_Spending,Current_Income,Current_Age,Retirement_Age,
                        Current_Checking,Current_Savings,Current_Brokerage,
                        Current_Roth_401k,Current_Pretax_401k,Current_Roth_IRA,Current_Pretax_IRA):
    INIT_SPENDING = Annual_Spending
    INIT_SALARY = Current_Income
    INIT_AGE = Current_Age
    INIT_RETIREAGE = Retirement_Age
    init_checkings = Current_Checking
    init_savings = Current_Savings
    init_brokerage = Current_Brokerage
    init_roth_401k = Current_Roth_401k
    init_pretax_401k = Current_Pretax_401k
    init_roth_IRA = Current_Roth_IRA
    init_pretax_IRA = Current_Pretax_IRA

    set_INIT_AGE(INIT_AGE)
                            
    def career_model(i):

        age = INIT_AGE + i
        if age < INIT_RETIREAGE:
            salary = INIT_SALARY*(1.02 + INFLATION_RATE)**i
        elif age < 65:
            salary = 0
        else:
            # social security
            salary = 12*2e3* (1.0 + INFLATION_RATE)**i

        # Hoping for a 0.25% increase in lifestype per year over inflation
        spending = INIT_SPENDING*(1.0025 + INFLATION_RATE)**i
        return spending, salary

    p = portfolio(
        retirement_age=INIT_RETIREAGE,
        init_checkings=init_checkings,
        init_savings=init_savings,
        init_brokerage=init_brokerage,
        init_roth_401k=init_roth_401k,
        init_pretax_401k=init_pretax_401k,
        init_roth_IRA=init_roth_IRA,
        init_pretax_IRA=0.0
    )

    ages = []
    spendings = []
    salaries = []
    nets = []
    message = ""
    do_i_win = True
    while p.net_worth() > 0 and p.age < 100:

        spending, salary = career_model(p.age - INIT_AGE)
        ages.append(p.age)
        spendings.append(spending)
        salaries.append(salary)
        nets.append(p.net_worth())
        try:
            msg = p.increment_year(pay=salary, spending=spending, verbose=False)
            if p.age % 5 == 0:
                message += "\n" + msg#print(msg)
        except:
            do_i_win = False
            if p.age % 5 == 0:
                message += "\nOUT OF MONEY, SO NO MORE LIVING"
                #print("OUT OF MONEY, SO NO MORE LIVING")
    fig = Figure()
    canvas = FigureCanvas(fig)
    ax = fig.gca()
    width, height = fig.get_size_inches() * fig.get_dpi()
    if do_i_win:
        ax.plot(ages, nets, 'g')
    else:
        ax.plot(ages, nets, 'r')
    ax.set_xlabel("Age")
    ax.set_ylabel("Net Worth")
    canvas.draw()
    image = np.frombuffer(canvas.tostring_rgb(),dtype='uint8').reshape(int(height),int(width), 3)
    return image,message

gr.Interface(fn=net_worth_calculator,
             inputs = [gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=31, maximum=100, value=30, step=1),
                       gr.Slider(minimum=31, maximum=100, value=65, step=1),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                       gr.Slider(minimum=0, maximum=1000e3, value=100e3, step=10e3),
                      ],
             outputs=['image','text']).launch()