import streamlit as st
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

st.title("Investment App 💼")

A = st.number_input(
    "Insert the initial investment (in $): "
)
r = st.number_input(
    "Insert the interest rate (nominal in %): "
)
T = st.number_input(
    "Insert the number of months of your investment: ",
    step=1
)

# Numpy array with the months
t = np.arange(T + 1)

# Numpy array for the interest rate
r_list = np.array([r] * len(t))

# Numpy array with the returns for each month in t
y = A * (1 + (r_list / (12*100)))**t

# Numpy array with the continuously compunded returns
y_c = A * np.exp(r_list * t / (12*100))

# Create the investment dictionary
inv_dict = {"Month": t, 
            "Rate": r_list, 
            "Returns": y, 
            "Continuous returns": y_c}

# Create the investment dataframe
inv_df = pd.DataFrame(inv_dict)
inv_df.columns = inv_dict.keys()

st.write(f"Table with returns at an interest rate of {r}")
new_df = st.data_editor(inv_df)

new_y = [A]
new_y_c = [A]

for i in range(1, len(t)):
    r_temp = new_df["Rate"].values[i]
    # New numpy array with the returns for each month in t
    val =  new_y[i-1] * (1 + r_temp / (12*100))
    new_y.append(val)

    # New numpy array with the continuously compunded returns
    val_c = new_y_c[i-1] * np.exp(r_temp / (12*100))
    new_y_c.append(val_c)

new_df["Returns"] = new_y
new_df["Continuous returns"] = new_y_c
new_df["Actual returns"] = 100 * (new_df["Continuous returns"] - new_df["Continuous returns"].shift(1)) / new_df["Continuous returns"].shift(1)

st.write("Table with variable rates")

st.dataframe(new_df)

fig, ax = plt.subplots()
ax.set_title("Investment returns plot")
ax.set_xlabel("Months")
ax.set_ylabel("Actual returns")
ax.plot(new_df["Month"], new_df["Actual returns"])

st.pyplot(fig)