# Import libraries
import datetime
import warnings

from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
from pandas_datareader import data as pdr

plt.style.use('bmh')
warnings.filterwarnings("ignore")

# Get historical data
startdate = datetime.datetime(1988, 1, 1)
enddate = datetime.datetime(2021, 1, 1)

# Pick SP 500 total return index
SP = pdr.get_data_yahoo('^SP500TR', start=startdate, end=enddate)
SP['date'] = SP.index

# Group by year
SP['year'] = SP['date'].apply(lambda d: f'{d.year}')
df = SP[['year', 'Adj Close']].groupby('year').first()
df['return'] = df['Adj Close'].pct_change() # return rate for last year's investment
df['return_next_y'] = df['return'].shift(-1) # return rate by the begining of next year
df = df.reset_index()
df = df.drop(columns='return')

# You can use the following group by dataframe if you want to do value averging monthly or quarterly
# # Group by month
# SP['month'] = SP['date'].apply(lambda d: f'{d.month:02d}')
# df_m = SP[['year', 'month', 'Adj Close']].groupby(['year','month']).first()
# df_m['return'] = df_m['Adj Close'].pct_change()
# df_m['return_next'] = df_m['return'].shift(-1)
# df_m = df_m.reset_index()
# df_m = df_m.drop(columns='return')

# # Group by quarter
# SP['quarter'] = SP['month'].apply(lambda m: int(m)//4 + 1)
# df_q = SP[['year', 'quarter', 'Adj Close']].groupby(['year','quarter']).first()
# df_q['return'] = df_q['Adj Close'].pct_change()
# df_q['return_next'] = df_q['return'].shift(-1)
# df_q = df_q.reset_index()
# df_q = df_q.drop(columns='return')

df.round(2).head(5)

def calc_value_averaging_return(df: pd.DataFrame, value_per_period: float = 1.0, 
                                increase_investment_per_period: float = 0.0,
                                display_irr: bool = True) -> pd.DataFrame:
    """
    This function calculates the internal rate of return of the value averaging strategy.
    It also returns a dataframe that contains the amount of money should invest or withdraw per period.
    :param df: Original dataframe that contains the price at the beginning of each period
    :param value_per_period: Investment value per period
    :param increase_investment_per_period: Increase the investment by x each period
    :param display_irr: Whether to display IRR results or not
    :return: A dataframe that contains all the relevant info
    """
    df_va = pd.DataFrame.copy(df)
    value = value_per_period
    df_va.loc[0,'target_value'] = value
    df_va.loc[0,'shares_should_have_in_total'] = value_per_period/df_va.loc[0,'Adj Close']
    df_va.loc[0,'shares_to_buy_or_sell'] = value_per_period/df_va.loc[0,'Adj Close']
    df_va.loc[0,'should_invest_or_withdraw'] = value
    df_va.loc[0,'total_invest_or_withdraw'] = value

    for i in range(1,len(df_va.index)):
        value += value_per_period*((1 + increase_investment_per_period)**i)
        df_va.loc[i,'target_value'] = value
        df_va.loc[i,'shares_should_have_in_total'] = df_va.loc[i,'target_value']/df_va.loc[i,'Adj Close']
        df_va.loc[i,'shares_to_buy_or_sell'] = df_va.loc[i,'shares_should_have_in_total'] - df_va.loc[i-1,'shares_should_have_in_total']
        df_va.loc[i,'should_invest_or_withdraw'] = df_va.loc[i,'shares_to_buy_or_sell']*df_va.loc[i,'Adj Close']
        df_va.loc[i,'total_invest_or_withdraw'] = df_va.loc[i-1,'total_invest_or_withdraw'] + df_va.loc[i,'should_invest_or_withdraw']

    if display_irr:
        cf = (-df_va['should_invest_or_withdraw']).tolist()
        cf[-1] = (df_va.loc[len(df_va.index)-2,'shares_should_have_in_total'] *
                    df_va.loc[len(df_va.index)-1,'Adj Close'])
        print("The cash flow (the negative sign stands for investment):")
        display(cf)
        irr = round(100*np.irr(cf), 2)
        print(f"The internal return rate is {irr}% per period")
        
    return df_va

def calc_buy_only_value_averaging_return(df: pd.DataFrame, value_per_period: float = 1.0, 
                                         increase_investment_per_period: float = 0.0,
                                         min_investment_per_period: float = 0.0,
                                         max_investment_per_period: float = 5.0,
                                         display_returns: bool = True) -> pd.DataFrame:
    """
    This function calculates the overall return rate and the internal return rate of the 
    buy only value averaging strategy, with a minimum and maximum investment amount per period.
    It also returns a dataframe that contains the amount of money should invest per peiod. 
    :param df: Original dataframe that contains the price at the beginning of each period
    :param value_per_period: Investment value per period
    :param increase_investment_per_period: Increase the investment by x each period
    :param min_investment_per_period: minimum investment amount at the beginning of each period
    :param max_investment_per_period: maximum investment amount at the beginning of each period
    :param display_returns: Whether to display return results or not
    :return: A dataframe that contains all the relevant info
    """
    df_va = pd.DataFrame.copy(df)
    value = value_per_period
    shares = value_per_period/df_va.loc[0, 'Adj Close']
    df_va.loc[0, 'target_value'] = value   # target value to reach at the beginning of each period
    df_va.loc[0, 'target_shares'] = shares # target share to have at the beginning of each period
    df_va.loc[0, 'current_shares'] = 0     # current shares before the period
    df_va.loc[0, 'shares_to_buy'] = shares # shares to add based on the target
    df_va.loc[0, 'total_shares'] = shares  # total shares after the investment at the beginning of each period
    df_va.loc[0, 'should_invest'] = value  # investment at the beginning of each period
    df_va.loc[0, 'total_invest'] = value   # total amount of investment including this period
    df_va.loc[0, 'total_value'] = value    # total share values after the investment at the beginning of each period
    df_va.loc[0, 'total_value_next'] = df_va.loc[0, 'total_value'] * (1 + df_va.loc[0, 'return_next_y'])
    df_va.loc[0, 'overall_return_next'] = df_va.loc[0, 'total_value_next']/df_va.loc[0, 'total_invest'] - 1
    
    for i in range(1,len(df_va.index)):
        df_va.loc[i, 'current_shares'] = df_va.loc[i-1, 'total_shares']
        
        value += value_per_period*((1 + increase_investment_per_period)**i)
        df_va.loc[i,'target_value'] = value
        df_va.loc[i,'target_shares'] = df_va.loc[i,'target_value']/df_va.loc[i,'Adj Close']
        df_va.loc[i,'shares_to_buy'] = df_va.loc[i,'target_shares'] - df_va.loc[i,'current_shares']
        df_va.loc[i,'should_invest'] = df_va.loc[i,'shares_to_buy'] * df_va.loc[i, 'Adj Close']
        
        if df_va.loc[i, 'should_invest'] < min_investment_per_period:
            df_va.loc[i, 'should_invest'] = min_investment_per_period
            df_va.loc[i, 'shares_to_buy'] = min_investment_per_period/df_va.loc[i,'Adj Close']
            
        elif df_va.loc[i, 'should_invest'] > max_investment_per_period:
            df_va.loc[i, 'should_invest'] = max_investment_per_period
            df_va.loc[i, 'shares_to_buy'] = max_investment_per_period/df_va.loc[i,'Adj Close']
        
        df_va.loc[i, 'total_shares'] = df_va.loc[i-1, 'total_shares'] + df_va.loc[i, 'shares_to_buy']
        df_va.loc[i, 'total_invest'] = df_va.loc[i-1, 'total_invest'] + df_va.loc[i, 'should_invest']
        
        df_va.loc[i, 'total_value'] = df_va.loc[i, 'total_shares']*df_va.loc[i, 'Adj Close']
        
        df_va.loc[i, 'total_value_next'] = df_va.loc[i, 'total_value'] * (1 + df_va.loc[i, 'return_next_y'])
    
        df_va.loc[i, 'overall_return_next'] = df_va.loc[i, 'total_value_next']/df_va.loc[i, 'total_invest'] - 1
        
    if display_returns:     

        # Calculate IRR
        cf = (-df_va['should_invest']).tolist()
        cf[-1] = df_va.loc[len(df_va.index) - 2, 'total_value_next']
        print("The cash flow (the negative sign stands for investment):")
        display(cf)
        irr = round(100*np.irr(cf), 2)
        print(f"The internal return rate is {irr}% per period")
              
        # Calculate overall return
        overall_return = round(df_va.loc[len(df_va.index) - 2, 'overall_return_next']*100, 2)
        print(f"The overall return rate (IRR) is {overall_return}%")
        
        # Calculate CAGR
        cagr = round(100*((df_va.loc[len(df_va.index) - 2, 'total_value_next']/
                           df_va.loc[len(df_va.index) - 2, 'total_invest'])**(1/(len(df_va.index) - 1)) - 1), 2)
        print(f"The compound annual growth rate (CAGR) is {cagr}% per period")
        
    return df_va

def calc_dollar_cost_averaging_return(df: pd.DataFrame, value_per_period: float = 1.0, 
                                      increase_investment_per_period: float = 0.0,
                                      display_returns: bool = True) -> pd.DataFrame:
    """
    This function calculates the overall rate of return and the internal rate of return of dollar cost averaging method. 
    :param df: Original dataframe that contains the price at the beginning of each period
    :param value_per_period: Investment value per period
    :param increase_investment_per_period: Increase the investment by x each period
    :param display_returns: Whether to display return results or not
    :return: A dataframe that contains all the relevant info
    """
    df_dc = pd.DataFrame.copy(df)
    # amount to invest each period
    df_dc.loc[0, 'should_invest'] = value_per_period   
    # total amount of investment after the investment this period
    df_dc.loc[0, 'total_invest'] = value_per_period   
    # the total value by the beginning of next period
    df_dc.loc[0, 'total_value_next'] = value_per_period*(1 + df_dc.loc[0, 'return_next_y'])
    df_dc.loc[0, 'overall_return_next'] = df_dc.loc[0, 'total_value_next'] / df_dc.loc[0, 'total_invest'] - 1

    for i in range(1,len(df_dc.index)):
        df_dc.loc[i, 'should_invest'] = value_per_period * (1 + increase_investment_per_period)**i
        df_dc.loc[i, 'total_invest'] = df_dc.loc[i-1, 'total_invest'] + df_dc.loc[i, 'should_invest']
        df_dc.loc[i, 'total_value_next'] = ((df_dc.loc[i-1, 'total_value_next'] + 
                                            df_dc.loc[i, 'should_invest']) *
                                            (1 + df_dc.loc[i, 'return_next_y']) 
                                           )
        df_dc.loc[i, 'overall_return_next'] = df_dc.loc[i, 'total_value_next'] / df_dc.loc[i, 'total_invest'] - 1
        
    if display_returns:
        
        # Calculate IRR
        cf = (-df_dc['should_invest']).tolist()
        cf[-1] = df_dc.loc[len(df_dc.index) - 2, 'total_value_next']
        print("The cash flow (the negative sign stands for investment):")
        display(cf)
        irr = round(100*np.irr(cf), 2)
        print(f"The internal return rate (IRR) is {irr}% per period")
                
        # Calculate overall return
        overall_return = round(df_dc.loc[len(df_dc.index) - 2, 'overall_return_next']*100, 2)
        print(f"The overall return rate is {overall_return}%")
        
        # Calculate CAGR
        cagr = round(100*((df_dc.loc[len(df_dc.index) - 2, 'total_value_next']/
                           df_dc.loc[len(df_dc.index) - 2, 'total_invest'])**(1/(len(df_dc.index) - 1)) - 1), 2)
        print(f"The compound annual growth rate (CAGR) is {cagr}% per period")
        
    return df_dc


def plot_changes(df: pd.DataFrame,
                 y1: str = 'total_invest', y2: str = 'total_value_next',
                 xlabel: str = 'period', ylabel: str = 'value',
                 title: str = 'Visualize total investment and value each year',
                 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
    :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()

df_va1 = calc_value_averaging_return(df, value_per_period=1200)

df_va1.round(2)

df_va2 = calc_value_averaging_return(df, value_per_period=1200, 
                                     increase_investment_per_period=0.1)

df_va2.round(2)

df_va3 = calc_buy_only_value_averaging_return(df, value_per_period=1200, 
                                             increase_investment_per_period=0.10,
                                             min_investment_per_period=600,
                                             max_investment_per_period=6000)

df_va3.round(2)

plot_changes(df_va3, title='Value averaging with adjustment: No selling, with min/max investment amount', xlabel='year')

# Dollar cost averaging: increase the investment per year by 2% 
df_dc = calc_dollar_cost_averaging_return(df, value_per_period=1200, 
                                          increase_investment_per_period=0.02)

df_dc.round(2)

# Dollar cost averaging: increase the investment per year by 2% 
plot_changes(df_dc, title='Yearly dollar cost averaging', xlabel='year')

# Compare dollar cost averaging(2% increase by year) with 
# value averaging method (no selling, with min and max value per period)
select_columns = ['should_invest', 'total_value_next', 'overall_return_next']
df_compare = pd.merge(df_dc[select_columns], df_va3[select_columns], 
         left_index=True, right_index=True, 
         suffixes=('_dc', '_va3')).round(2)
df_compare

def compare_strategy(df, str1, str2, str1_name, str2_name):
    print(f'number of times {str1_name} is better than {str2_name}: {df.loc[df[str1] > df[str2],:].shape[0]}')
    print(f'number of times {str1_name} is equal to {str2_name}: {df.loc[df[str1] == df[str2],:].shape[0]}')
    print(f'number of times {str1_name} is worse than {str2_name}: {df.loc[df[str1] < df[str2],:].shape[0]}')

compare_strategy(df_compare, 'overall_return_next_dc', 'overall_return_next_va3', 
                 'yearly dollar-cost-averaging', 'adjusted version of value averaging')