# Import libraries
from matplotlib import pyplot as plt
import pandas as pd

plt.style.use('bmh')

def calc_return_over_periods(initial_investment: int, investment_per_period: int,
                             return_per_period: float, n_period: int,
                             increase_investment_per_period: float = 0.00,
                             invest_at_period_begin: bool = True) -> pd.DataFrame:
    """
    This function calculates the overall rate of return across n periods given a certain return rate pre period
    :param initial_investment: Amount of initial investment
    :param investment_per_period: Amount of investment per period
    :param return_per_period: Assumed rate of return per period
    :param n_period: Total number of periods
    :param increase_investment_per_period: Increase the amount of investment by x for each period after the first investment period
    :param invest_at_period_begin: whether to invest at the begining of each period or not
    :return: dataframe contains the rate of return for n periods.
    """
    df_result = pd.DataFrame()

    for i in range(n_period):
        if i == 0:
            if invest_at_period_begin:
                period_begin_investment = initial_investment + investment_per_period
            if not invest_at_period_begin:
                period_begin_investment = initial_investment
            df_result.loc[i, 'period_begin_investment'] = period_begin_investment
            df_result.loc[i, 'total_investment'] = period_begin_investment
            value = period_begin_investment * (1 + return_per_period)
            df_result.loc[i, 'value_by_next_period'] = value
        else:
            if invest_at_period_begin:
                period_begin_investment = investment_per_period * ((1 + increase_investment_per_period) ** i)
            if not invest_at_period_begin:
                period_begin_investment = investment_per_period * ((1 + increase_investment_per_period) ** (i - 1))
            value += period_begin_investment
            df_result.loc[i, 'period_begin_investment'] = period_begin_investment
            df_result.loc[i, 'total_investment'] = (df_result.loc[i - 1, 'total_investment'] +
                                                    df_result.loc[i, 'period_begin_investment'])
            value = value * (1 + return_per_period)
            df_result.loc[i, 'value_by_next_period'] = value

    df_result['return_by_next_period'] = df_result['value_by_next_period'] / df_result['total_investment'] - 1

    return df_result.round(2)


def plot_changes(df: pd.DataFrame,
                 y1: str = 'total_investment', y2: str = 'value_by_next_period',
                 xlabel: str = 'period', ylabel: str = 'value',
                 title: str = 'Visualize your saving plan and its future value',
                 target: int = None) -> None:
    """
    This function visualizes the total investment and the total value of the investment plan across time
    :param df: A dataframe contains total investment and total value
    :param y1: Column name
    :param y2: Column name
    :param xlabel: X label
    :param ylabel: Y label
    :param title: Title of the graph
    :param target: Final value one wants to achieve
    :return: A plot
    """
    plt.figure(figsize=(15, 5))
    plt.plot(df[y1], label=y1)
    plt.plot(df[y2], label=y2)
    if target:
        plt.axhline(y=target, c='black', label='target value')
    plt.xlim([0, len(df[y1])-1])
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.title(title)
    plt.legend()
    plt.show()

# Each period is a year

# Assuming investing 20000 Euro initially 
initial_investment = 20000

# Assuming investing 10000 Euro in the begining of each year
investment_per_period = 10000

# Assuming the average investment return rate is about 5%
return_per_period = 0.05

# Assuming the plan lasts for 30 years
n_period = 30

df_plan1 = calc_return_over_periods(initial_investment, investment_per_period, 
                                   return_per_period, n_period)
df_plan1

target_value = 720000

plot_changes(df_plan1, title='Case 1', target = target_value, xlabel='year')

# Each period is a year

# Assuming investing 20000 Euro initially 
initial_investment = 20000

# Assuming investing 10000 Euro in the begining of each year
investment_per_period = 10000

# Assuming the average investment return rate is about 5%
return_per_period = 0.05

# Assuming the plan lasts for 30 years
n_period = 30

# Assuming increase the amount of investment each year by 2%
increase_investment_per_period = 0.02

df_plan2 = calc_return_over_periods(initial_investment, investment_per_period, 
                                   return_per_period, n_period,
                                   increase_investment_per_period)

df_plan2

plot_changes(df_plan2, title='Case 2', target=target_value, xlabel='year')

# Each period is a month

# No initial investment
initial_investment = 0

# Invest 500 Euro at the end of each month
investment_per_period = 500

# Assuming the average investment return rate is about 5% per year
return_per_year = 0.05
return_per_period = (return_per_year + 1)**(1/12) - 1 

# Assuming the plan lasts for 40 years
n_years = 40
n_period = n_years*12

df_plan3 = calc_return_over_periods(initial_investment, investment_per_period, 
                                   return_per_period, n_period,
                                   invest_at_period_begin=False)

df_plan3

plot_changes(df_plan3, title='Case 3', target=target_value, xlabel='month')