import statsmodels.api as sm
import pandas as pd
from dateutil.relativedelta import relativedelta
from QuantConnect.Indicators import SimpleMovingAverage
from functools import partial 

class OptimizedNadionRadiator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 6)  # Set Start Date
        self.SetEndDate(2020, 1, 14)
        self.SetCash(100000)  # Set Strategy Cash
        
        symbols = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.AddAlpha(MyAlphaModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())


class MyAlphaModel(AlphaModel):
    symbol = None
    sma = None
    model = None
    
    def Update(self, algorithm, slice):
        # Only emit insights if the exchange is on its last minute of the day (proxy for daily close)
        exchange = algorithm.Securities[self.symbol].Exchange
        if not (exchange.DateTimeIsOpen(slice.Time) and not exchange.DateTimeIsOpen(slice.Time + relativedelta(minutes=1))):
            return []
        
        close = slice[self.symbol].Close
        
        self.sma.Update(slice.Time, slice[self.symbol].Close)
        normalized_sma = (close - self.sma.Current.Value) / close
        normalized_sma = pd.DataFrame({"const": [1], "normalized_SMA": [normalized_sma]})
        prediction = self.model.predict(normalized_sma).values[0]

        if prediction == 0:
            return []
        
        direction = InsightDirection.Up if prediction > 0 else InsightDirection.Down
        return [Insight.Price(self.symbol, timedelta(days=1), direction)]
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            self.symbol = security.Symbol
            self.sma = SimpleMovingAverage(5)

            history = algorithm.History(self.symbol, 50, Resolution.Daily).loc[self.symbol][['close']]
            for idx, row in history.iterrows():
                self.sma.Update(idx, row.close)
                if self.sma.IsReady:
                    history.loc[idx, "SMA"] = self.sma.Current.Value
            
            history['normalized_SMA'] = (history.close - history.SMA) / history.close
            history['return'] = history.close.pct_change()
            history.dropna(inplace=True)
            
            X = sm.add_constant(history.normalized_SMA)
            y = history['return'].values
            self.model = sm.OLS(y, X).fit()