# Yfin_S&P+PowLaw
# Analysis of S&P500 index
# yfinance https://ranaroussi.github.io/yfinance/
# Downloads S&P500 index '^GSPC'
# Fits compound annual growth power law function a*(1+CAGR)^b for annual returns
# Calculates index - power law difference to determine extent of deviations and mean reversions
# Calculates autocorrelation function to highlight sjourn times of economic cycles 

import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import matplotlib.dates as mdates
from scipy.optimize import curve_fit
from scipy import signal
from scipy.fft import fft, fftfreq
import numpy as np

def download_sp500_data():
    # Define the ticker symbol for S&P 500
    ticker_symbol = '^GSPC' #^IXIC'

    # Define the start and end dates for the data
    start_date = "1928-01-03"
    end_date = "2025-10-10"

    # Download historical data from Yahoo Finance
    sp500_data = yf.download(ticker_symbol, start=start_date, end=end_date, auto_adjust=True)
    
    return sp500_data


def power_law_with_baseline(x, a, b): #+):
    return a*(b+0.03)**(x/365)

    
def fit_power_law_with_baseline(x, y):
    # Ensure x and y are NumPy arrays and 1D
    x = np.asarray(x, dtype=float).flatten()
    y = np.asarray(y, dtype=float).flatten()

    print("Fitting data: X shape:", x.shape, "Y shape:", y.shape)  # Debugging output

    # Fit a power-law with a baseline to the data
    try:
        params, _ = curve_fit(power_law_with_baseline, x, y, maxfev=5000)
        return params
    except Exception as e:
        print("Error during curve fitting:", str(e))
        return [1, 1, 1]  # Return default parameters if fitting fails
        

def calculate_power_law_ratio(data):
    if 'Close' not in data:
        raise ValueError("Data does not contain 'Close' column.")

    # Generate x values (indices)
    x_values = np.arange(len(data))

    # Ensure Close prices are numeric and drop NaN values
    #data = data.dropna(axis=['Close']) #subset=
    y_values = data['Close'].values.astype(float)
    
    # Fit the power-law model
    params = fit_power_law_with_baseline(x_values[:len(y_values)], y_values)
    
    # Calculate the fitted values using the power-law with baseline
    #data['PowerLaw_Fit'] = power_law_with_baseline(x_values[:len(y_values)], *params)
    data['PowerLaw_Fit'] = power_law_with_baseline(x_values[:len(y_values)], *params)

    # Calculate the ratio of Close price to the power-law fit
    data['Ratio'] = data['Close'].div(data['PowerLaw_Fit'], axis=0)
    data['Difference'] = data['Close'].sub(data['PowerLaw_Fit'], axis=0)
    #data['Long_MA'] = data['Close'].rolling(window=200, min_periods=1).mean()

    corr = signal.correlate(data['Difference'], data['Difference'], mode="full")
    lags = signal.correlation_lags(len(data['Difference']), len(data['Difference']))
    corr /= np.max(corr)
    data['Lags'] = lags[lags.size//2:]
    data['Corr'] = corr[corr.size//2:]

    dotcom_close = data.loc['2000-07-03', 'Close'].values.astype(float)
    dotcom_loss = data.loc['2000-07-03', 'Difference'].values

    oct2025_close = data.loc['2025-09-10', 'Close'].values.astype(float)
    oct2025_loss = data.loc['2025-09-10', 'Difference'].values
    
    print(f'Dot com peak date: 2000-07-03 close: {dotcom_close}, reversion: {-dotcom_loss}, pc from top: {-100*dotcom_loss/dotcom_close}')
    print(f'Oct 2025 high data: 2025-09-10 close: {oct2025_close}, reversion: {-oct2025_loss}, pc from top: {-100*oct2025_loss/oct2025_close}')
    
    return data, params
    

def plot_data(data, params):

    plt.rcParams['font.size'] = 9
    # Close price
    a, b = params
    a = np.round(a, decimals = 2)
    b = np.round(b, decimals = 2)
    cagr = np.round((b-1)*100, decimals=0)
    plt.figure(figsize=(8,8))
    ax1 = plt.subplot(311)
    ax1.plot(data['Close'], label='Close Price', color='blue', lw = 1.5)
    ax1.plot(data['PowerLaw_Fit'], label=f'Power-Law a*b$^x$: a: {str(a)} b: {str(b)} CAGR = {str(cagr)}%', color='red', linestyle='--', lw = 1.5)
    #plt.plot(data.index, data['Long_MA'], label='200 day MA', color = 'green', linestyle='--')
    ax1.set_xlim(np.datetime64('1927-12-30'), np.datetime64('2030-10-03'))
    ax1.xaxis.set_major_locator(MultipleLocator(10*365.25))
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
    #ax1.tick_params(axis='x', labelrotation=45)
    ax1.xaxis.set_minor_locator(MultipleLocator(365))
    #ax1.legend()
    #plt.yscale('log')
    plt.legend(loc='lower left', bbox_to_anchor=(0, 0.1), fontsize = 9, frameon=False)
    #plt.title('S&P 500 Price and Power-Law Fit')
    #plt.show()

    axins1 = inset_axes(ax1, width="63%", height="50%", loc="upper left")
    axins1.plot(data['Close'], label='Close Price', color='royalblue', lw = 1.5)
    axins1.plot(data['PowerLaw_Fit']+12.5, color='red', linestyle='--', lw = 1.5)
    #plt.plot(data.index, data['Long_MA'], label='200 day MA', color = 'green', linestyle='--')
    axins1.set_xlim(np.datetime64('1927-12-30'), np.datetime64('1993-12-30'))
    axins1.xaxis.set_major_locator(MultipleLocator(10*365.25))
    axins1.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
    axins1.tick_params(axis='x', labelrotation=0)
    axins1.set_xticklabels([])
    #axins1.xaxis.set_minor_locator(MultipleLocator(365))
    axins1.set_ylim(0, 500)
    axins1.yaxis.set_label_position("right")
    axins1.yaxis.tick_right()
    axins1.spines[['top', 'left']].set_visible(False)
    
    # Plot Close - CAGR
    #plt.figure(figsize=(8, 4))
    ax2 = plt.subplot(312)
    plt.plot(data.index, data['Difference'], label=f'Close - Power', color='darkorange', lw = 1.5)
    plt.axhline(1, color='gray', linestyle='--', label='Baseline (Diff = 0)')
    ax2.set_xlim(np.datetime64('1927-12-30'), np.datetime64('2030-10-03'))
    ax2.xaxis.set_major_locator(MultipleLocator(10*365.25))
    ax2.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
    ax2.xaxis.set_minor_locator(MultipleLocator(365))
    ax2.set_ylim(-1000,1500)
    plt.legend(loc='lower left', bbox_to_anchor=(0, 0), frameon=False)
    #plt.title('Close Price - Power-Law')
    #plt.show()

    axins2 = inset_axes(ax2, width="63%", height="50%", loc="upper left")
    axins2.plot(data['Difference']-12.5, label=f'Close - Power', color='orange', lw = 1.5)
    axins2.axhline(1, color='gray', linestyle='--')
    #plt.plot(data.index, data['Long_MA'], label='200 day MA', color = 'green', linestyle='--')
    axins2.set_xlim(np.datetime64('1927-12-30'), np.datetime64('1993-12-30'))
    axins2.xaxis.set_major_locator(MultipleLocator(10*365.25))
    axins2.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
    axins2.tick_params(axis='x', labelrotation=0)
    #axins2.xaxis.set_minor_locator(MultipleLocator(365))
    #axins2.set_xticks([])
    # moving bottom spine up to y=0 position:
    axins2.xaxis.set_ticks_position('bottom')
    axins2.spines['bottom'].set_position(('data',1.25))
    axins2.set_xticklabels([])
    axins2.set_ylim(-70, 120)
    axins2.yaxis.set_label_position("right")
    axins2.yaxis.tick_right()
    axins2.spines[['bottom', 'top', 'left']].set_visible(False)

    # Plot the autocorrelation
    plt.subplot(313)
    plt.plot(data['Lags']/365, data['Corr'], label='AutoCorr(Close-Power)', color='green')
    plt.axhline(0, color='gray', linestyle='--')
    plt.legend(loc='upper right', bbox_to_anchor=(1, 1), fontsize = 9, frameon=False)
    plt.xlim(0,24)
    plt.xticks(np.arange(0,24,2))
    #output_path = ''
    #plt.savefig(output_path+"\\Yfin_S&PAnalysis.png", bbox_inches = 'tight', transparent=False)
    plt.show()
    

def main():
    sp500_data = download_sp500_data()
    sp500_data_with_ratio, params = calculate_power_law_ratio(sp500_data)
    plot_data(sp500_data_with_ratio, params)

if __name__ == "__main__":
    main()