import pandas as pd
import numpy as np
import yfinance as yf
from datetime import timedelta, datetime
import matplotlib.pyplot as plt

# Download S&P 500 historical data
df = yf.download('^SPX', start='1900-01-01')
spy = df[['Close']].resample('D').ffill().rename(columns={'Close': 'price'})
spy = spy.resample('MS').last()

# Parameters
MONTHLY_INVESTMENT = 500
HORIZONS = [5, 7, 10, 15, 20, 25, 35]  # Different investment horizons in years
dates = spy.index

def create_fanchart(arr, ax, color='blue'):
    x = np.arange(arr.shape[0])
    # for the median use `np.median` and change the legend below
    mean = np.mean(arr, axis=1)
    offsets = (10, 20, 30, 40)

    ax.plot(mean, color='black', lw=2)
    for offset in offsets:
        low = np.percentile(arr, 50-offset, axis=1)
        high = np.percentile(arr, 50+offset, axis=1)
        # since `offset` will never be bigger than 50, do 55-offset so that
        # even for the whole range of the graph the fanchart is visible
        alpha = (55 - offset) / 100
        ax.fill_between(x, low, high, color=color, alpha=alpha)
    ax.legend(['Mean'] + [f'Pct{2*o}' for o in offsets])
    return ax

# Function to simulate DCA strategy for different horizons
def sim_dca(horizon, start_date):
    dates_in_sim = dates[(dates >= start_date) & (dates < start_date + timedelta(horizon * 365))]
    investments = [{'date': date, 
                    'money_spent': MONTHLY_INVESTMENT, 
                    'quantity_bought': MONTHLY_INVESTMENT / spy.loc[date].price} 
                   for date in dates_in_sim]

    invest_df = pd.DataFrame(investments).set_index('date')
    invest_df['tt_qty'] = invest_df['quantity_bought'].cumsum()
    invest_df['index_price'] = spy['price'].reindex(invest_df.index)
    invest_df['market_value'] = invest_df['index_price'] * invest_df['tt_qty']
    invest_df['total_money_spent'] = invest_df['money_spent'].cumsum()
    
    
    return invest_df

# Function to calculate statistics for different horizons
def calculate_statistics(horizon):
    mc = []
    dca_results = []
    lumsum_results = []
    for start_date in dates[dates < dates.max() - timedelta(horizon * 365)]:
        invest_df = sim_dca(horizon, start_date)
        benefit_strategy = (invest_df['market_value'].iloc[-1] - invest_df['total_money_spent'].iloc[-1]) / invest_df['total_money_spent'].iloc[-1]
        lumpsum_benefit = (invest_df['total_money_spent'].iloc[-1] * 
                           (invest_df['index_price'].iloc[-1] / invest_df['index_price'].iloc[0]) - 
                           invest_df['total_money_spent'].iloc[-1]) / invest_df['total_money_spent'].iloc[-1]
        
        
        mc.append({'start_date': start_date, 
                   'benefit_strategy': benefit_strategy, 
                   'lumpsum_benefit': lumpsum_benefit})
        
        lumsum_results.append((spy.loc[invest_df['market_value'].index[0]: invest_df['market_value'].index[-1]].reset_index(drop=True)['price'].pct_change().fillna(0)+1).cumprod())
        dca_results.append(((invest_df['market_value'] / invest_df['total_money_spent']).reset_index(drop=True).pct_change().fillna(0)+1).cumprod())

    # Create subplots
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(5, 6))  # 2 rows, 1 column
    
    # DCA results
    ax1 = create_fanchart(np.array(dca_results).T, ax1, color='blue')
    ax1.set_title(f"DCA Return for period {horizon} years")
    
    # Lumsum results
    ax2 = create_fanchart(np.array(lumsum_results).T, ax2, color='green')
    ax2.set_title(f'Lumpsum Return for period {horizon} years')
    
    plt.tight_layout()  # Adjust layout to prevent overlap
    plt.show()

    return pd.DataFrame(mc).set_index('start_date')

for horizon in HORIZONS:
    mc_df = calculate_statistics(horizon)
    stats = mc_df['benefit_strategy'].describe().to_dict()
    
    print(f"\n### Horizon: {horizon} Years ###")
    print(f'Average percentage gain using DCA: {stats["mean"]*100:.1f}%')
    print(f'Median percentage gain using DCA: {stats["50%"]*100:.1f}%')
    print(f'q25 percentage gain using DCA: {stats["25%"]*100:.1f}%')
    print(f'q75 percentage gain using DCA: {stats["75%"]*100:.1f}%')
    
    proba_of_losing_dca = len(mc_df[mc_df['benefit_strategy'] < 0]) / len(mc_df)
    print(f'Probability of losing money with DCA: {proba_of_losing_dca*100:.1f}%')
    
    stats = mc_df['lumpsum_benefit'].describe().to_dict()
    print(f'Average percentage gain using lumpsum: {stats["mean"]*100:.1f}%')
    print(f'Median percentage gain using lumpsum: {stats["50%"]*100:.1f}%')
    print(f'q25 percentage gain using lumpsum: {stats["25%"]*100:.1f}%')
    print(f'q75 percentage gain using lumpsum: {stats["75%"]*100:.1f}%')
    
    proba_of_losing_lumpsum = len(mc_df[mc_df['lumpsum_benefit'] < 0]) / len(mc_df)
    print(f'Probability of losing money with lumpsum: {proba_of_losing_lumpsum*100:.1f}%')