import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import adfuller

def TestStationartiy(returns):
    
    # Return pandas Series with True/False for each symbol
    return pd.Series([adfuller(values)[1] < 0.05 for columns, values in returns.iteritems()], index = returns.columns)
    
def GetZScores(returns):
    # Return pandas DataFrame containing z-scores
    return returns.subtract(returns.mean()).div(returns.std())

from HMM import *
import numpy as np
from StationarityAndZScores import *
from QuantConnect.Data.UniverseSelection import *
from collections import deque
from System.Drawing import Color

class ModulatedDynamicCircuit(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.models = {}
        self.queue = deque(maxlen=2)
        self.SetWarmUp(2)
        
        self.AddEquity('SPY', Resolution.Daily)
        #self.Schedule.On(self.DateRules.EveryDay('SPY'), self.TimeRules.AfterMarketOpen('SPY', 5), self.GenerateInsights)
        #self.Schedule.On(self.DateRules.MonthStart('SPY'), self.TimeRules.At(19,0), self.RefitModels)

        stockPlot = Chart('HMM Plots')
        stockPlot.AddSeries(Series('HMM0', SeriesType.Scatter, '$', Color.Red, ScatterMarkerSymbol.Circle))
        stockPlot.AddSeries(Series('HMM1', SeriesType.Scatter, '$', Color.Blue, ScatterMarkerSymbol.Circle))
        self.AddChart(stockPlot)


    def RefitModels(self):
        for symbol, model in self.models.items():
            RefitModel(self, symbol, model)

    
    def OnData(self, data):

        self.queue.append(data["SPY"].Close)
        
        if self.IsWarmingUp:
            return
        
        symbol = "SPY"
        model = self.CheckForHMM(symbol)
        state_prediction = PredictState(self, model, symbol)

        self.Log("State = " + str(state_prediction))

        if state_prediction == 1:
            self.Plot('HMM Plots', 'HMM1', self.CurrentSlice[symbol].Close)
        elif state_prediction == 0:
            self.Plot('HMM Plots', 'HMM0', self.CurrentSlice[symbol].Close)
            self.Log(state_prediction)

    
    def CheckForHMM(self, symbol):
        if self.models.get(symbol, None) is None:
            self.models[symbol] = CreateHMM(self, symbol)
        return self.models.get(symbol, None)

    
    def OnSecuritiesChanged(self, changes):
        symbols = [x.Symbol for x in changes.AddedSecurities]
        # Build model for each symbol
        for symbol in symbols:
            self.models[symbol] = CreateHMM(self, symbol)

import numpy as np
import pandas as pd
from hmmlearn.hmm import GaussianHMM

def CreateHMM(algorithm, symbol):
    start_time = datetime(2010, 1, 1) # start datetime for history call
    end_time = datetime(2015, 1, 1) # end datetime for history call

    history = algorithm.History([symbol], start_time, end_time, Resolution.Daily)
    returns = np.array(history.loc[symbol].close.pct_change().dropna())
    # Reshape returns
    returns = np.array(returns).reshape((len(returns),1))
    # Initialize Gaussian Hidden Markov Model
    return GaussianHMM(n_components=2, covariance_type="full", n_iter=1000).fit(returns)
    
def PredictState(algorithm, model, symbol):
    returns = np.array(algorithm.queue)
    returns = returns[-2:]
    returns = pd.Series(returns)
    returns = np.array(returns.pct_change().dropna())
    returns = np.array(returns).reshape((len(returns),1))
    state = model.predict(returns)[0]
    return state
    
def RefitModel(algorithm, symbol, model):
    history = algorithm.History([symbol], 900, Resolution.Daily)
    returns = np.array(history.loc[symbol].close.pct_change().dropna())
    # Reshape returns
    returns = np.array(returns).reshape((len(returns),1))
    return model.fit(returns)