# Created by: Sergiu Iatco / 2021.10.01
# https://github.com/itsergiu/Predict-S-P-500-correction-with-Shiller-PE-Ratio
# How to predict S&P 500 correction

# Nadeem Walayat - The Market Oracle
# http://www.marketoracle.co.uk/Article69423.html

# Vitaliy Katsenelson 
# https://contrarianedge.com/sideways-market/

# https://www.marketwatch.com/story/market-analysts-cant-agree-on-where-stocks-are-going-next-so-double-check-the-data-before-you-buy-or-sell-11632447577

# Purpose: to build a machine learning model to predict S&P 500 correction within next 6 months

# !pip install pandas
# !pip install dateparser
# !pip install xgboost
# !pip install sklearn
# !pip install numpy
# !pip install plotly_express
# !pip install fredapi

import pandas as pd
# import dateparser
import xgboost as xgb
from sklearn import model_selection
from sklearn.metrics import r2_score
import numpy as np
from fredapi import Fred
import plotly.express as px
from IPython.display import display

url_per='https://www.multpl.com/shiller-pe/table/by-month'

ls_tables = pd.read_html(url_per)

df_per=ls_tables[0]

df_per.rename(
    columns=({ cols[1]: 'PER'}), 
    inplace=True,)

url_sp='https://www.multpl.com/s-p-500-historical-prices/table/by-month'

ls_tables = pd.read_html(url_sp)

df_sp.rename(
    columns=({ cols[1]: 'Price'}), 
    inplace=True,)

for i,e in enumerate(ls_price_p):
    df[e]=df.Price.shift(i)   

df['Price1P'] = df['Price'].shift(-1)
df['Price_Var']=df['Price']/df['Price1P']-1

df['PriceFmin']=df[ls_price_p].min(axis=1)

col_y = 'Price_Corr_6M'
df[col_y]=1-df['PriceFmin']/df['Price0F']

df_initial=df.copy()

df.dropna(axis=0, how='any', inplace=True) # remove NaN rowss

# plot PER & Price_correction

min_correction=0.05 # force learning above ?
# min_correction=0
if min_correction!=0:
    df[col_y] = df[col_y].apply(lambda x: 0 if x<min_correction else x)

# df[col_y] = -df[col_y] # reverse sign

df_per_pc=df[['Date','PER',col_y]].copy()
df_per_pc[[col_y]]=100*df_per_pc[[col_y]]
df_per_pc[col_y]=-df_per_pc[col_y] # reverse sign for plotting

def gen_features(df,cols, ref_period, band):
    features = []
    
    features.append(df['Date'].iloc[ref_period])
    for col in cols:
        segment = df[col].iloc[ref_period:ref_period+band]
        features.append(segment.min()) 
        features.append(segment.max())

        for i in range(2,99,10):
            features.append(np.quantile(segment, i/100))

        features.append(segment.mean())
        features.append(segment.std())
        features.append(segment.mad())
        features.append(segment.skew())
        features.append(segment.kurtosis())

        features.append(np.sqrt(np.mean(segment**2)))

        features.append(np.abs(segment).mean())
        features.append(np.abs(segment).std())

        features.append(np.abs(np.diff(segment)).mean())
        features.append(np.abs(np.diff(segment)).std())
        
    return pd.Series(features)

def generate_dataset(df_f, cols, band=60, start_date='2000-01-01', periods = 0):
#     global cols
    df = df_f.copy()
    X_train = pd.DataFrame()

    y_train = pd.Series(dtype='float64', name=col_y)
#     cols =['PER', 'Price', 'Price_Var']

    if periods>0:
        ref_periods = periods
    else:
        ref_periods = df[df.index>=start_date].shape[0]
        
    for ref_period in range(ref_periods):
        X = gen_features(df,cols, ref_period, band)
        X_train = X_train.append(X, ignore_index=True)
        y_train = y_train.append(pd.Series(df[col_y][ref_period], index=[df.index[ref_period]], name=col_y))

    X_train = X_train.rename({0:'Date'}, axis='columns')
    X_train = X_train.set_index('Date')
    return X_train, y_train

start_date = '1950-01-01' # set your own date
band=60 # how many months from past to use to generate statistics

cols_features = ['PER', 'Price', 'Price_Var']
X_train, y_train = generate_dataset(df, cols_features, band, start_date)

Y=dataset.pop(col_y) # target to predict

# PREDICTION

def generate_model(df, cols, start_date, verbosity = False):
    best_score = 0
    best_model = None
    best_band = 0
    test_size = 0.3
    bands = 60
    
    for band in range(12,bands+1,12): # [12,24,36,48,60] months
        X_train, y_train = generate_dataset(df, cols, band, start_date)
        dataset = pd.concat([y_train, X_train], axis=1, join="inner") # align to avoid random shifting
        Y=dataset.pop(col_y)
        X=dataset
        X_train, X_test, Y_train, Y_test = model_selection.train_test_split(X, Y, test_size=test_size, shuffle = True)
        model = xgb.XGBRegressor()
        model.fit(X_train,Y_train.values)
        score = model.score(X_test, Y_test)
        if best_score < score:
            best_score = score
            best_model = model
            best_band = band
            X_t = X_test
            Y_t = Y_test
            
#         if verbosity: # verbosit for each band
#             print(score, band)
    if verbosity: # verbosity for best score
        print('retest score:', best_model.score(X_t, Y_t).round(3), best_score.round(3), 'features:', cols, 'X_train.shape:', X_train.shape)# retest best model score
    
    return best_model, best_band, best_score

# model.score(X_t, Y_t), best_score # retest best model score

# df_pred = df_pred.iloc[:periods]
# df_pred

y_pred = model.predict(X_pred).round(2) # readable
# y_pred = model.predict(X_pred)
# y_pred

df_pred = df_pred.iloc[:periods].copy()
df_pred[col_y+'Pred'] = np.clip(y_pred,0,0.5) # replace outliers