import numpy as np
import pandas as pd
from sklearn.neural_network import MLPClassifier

class MachineLearningSPY(QCAlgorithm):
    
    def Initialize(self):
        
        self.SetStartDate(2016, 5, 2)  
        self.SetEndDate(2021, 6, 22)  
        self.SetCash(100000)  
        self.AddEquity("SPY", Resolution.Daily)  
        self.SetBenchmark("SPY")
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.AddUniverseSelection(ManualUniverseSelectionModel("SPY"))
        
        self.lookback = 30
        self.SetWarmup(self.lookback)
        self.spyClose = RollingWindow[float](self.lookback)
        self.spyMomentum = RollingWindow[float](self.lookback)
        self.spyMomentum_indicator = self.MOMP("SPY", self.lookback, Resolution.Daily)
        self.AddAlpha(MachineLearningSPYAlphaModel(self.spyClose, self.spyMomentum, self.spyMomentum_indicator))
        
class MachineLearningSPYAlphaModel:
    def __init__(self, close, spyMomentum, spyMomentum_indicator):
        self.period = timedelta(30)
        self.spyClose = close
        self.spyMomentum = spyMomentum
        self.spyMomentum_indicator = spyMomentum_indicator
    
    def GetMLModel(self):
        self.MLModel = 0
        self.MLModel = MLPClassifier(hidden_layer_sizes = (100, 100, 100, 100), max_iter = 1000)
        
    def Update(self, algorithm, data):
        insights = []
        
        if data.Bars.ContainsKey("SPY"):
            self.spyMomentum_indicator.Update(data["SPY"].EndTime, data["SPY"].Close)
            self.spyMomentum.Add(self.spyMomentum_indicator.Current.Value)
            self.spyClose.Add(data["SPY"].Close)
            
            if not algorithm.IsWarmingUp and self.spyMomentum.IsReady and self.spyClose.IsReady:
                
                # features dataframe
                df1 = pd.DataFrame(self.spyMomentum, columns=["MOM"]).reset_index(drop=True)
                df2 = pd.DataFrame(self.spyClose, columns=["Close"]).reset_index(drop=True)
                self.df = pd.concat([df1, df2], axis=1)
                
                # calculate daily SPY forward returns to be used to set Target / Signal
                self.df['spyReturn'] = np.log(self.df['Close'].shift(-1)/self.df['Close']) 
                self.df = self.df.dropna()
                
                # set Signal / Target
                self.df['Signal'] = 0
                self.df.loc[self.df['spyReturn'] > 0, 'Signal'] = 1
                self.df.loc[self.df['spyReturn'] < 0, 'Signal'] = -1
                
                # set training data
                self.X = self.df.drop(['Close','Signal'], axis=1)
                self.Y = self.df['Signal']
        
                self.Y, self.X = self.Y.align(self.X, axis=0, join='inner')
                self.X_train = self.X[:-1]
                self.Y_train = self.Y[:-1]
                self.X_train.replace([np.inf, -np.inf], np.nan, inplace=True)
                self.Y_train.replace([np.inf, -np.inf], np.nan, inplace=True)
                
                drops = []
                [drops.append(i) for i in range(self.X_train.shape[0]) if self.X_train.iloc[i].isnull().any()]
                [drops.append(i) for i in range(self.Y_train.shape[0]) if self.Y_train.iloc[i] == np.nan and i not in drops]
                self.X_train.drop(index=self.X_train.index[drops], inplace=True)
                self.Y_train.drop(index=self.Y_train.index[drops], inplace=True)
                if self.X_train.empty or self.Y_train.empty: return []
                
                # fit / train ML model
                self.GetMLModel()
                self.MLModel.fit(self.X_train, self.Y_train)
                
                # predict next day signal using today's values of feature set
                self.X_today = self.X.iloc[-1]
                # self.X_today is Series, so convert to numpy array
                self.X_today = self.X_today.to_numpy()
                # reshape self.X_today because it only has 1 day's sample
                self.X_today = self.X_today.reshape(1,-1)
                
                # Y_predict will take predicted signal
                self.Y_predict = self.Y.iloc[-1]
                try:
                    self.Y_predict = self.MLModel.predict(self.X_today)
                except: return []
                
                if self.Y_predict == 1:
                    insights.append(Insight("SPY", self.period, InsightType.Price, InsightDirection.Up, 1, None))
                elif self.Y_predict == -1:
                    insights.append(Insight("SPY", self.period, InsightType.Price, InsightDirection.Down, 1, None))
                    
        return insights
                
    def OnSecuritiesChanged(self, algorithm, changes):
        self.changes = changes