9#region imports
from AlgorithmImports import *
from statsmodels.tsa.regime_switching.markov_regression import MarkovRegression
#endregion

class HMMDemo(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2019, 6, 1)
        self.SetEndDate(2024, 4, 1)
        self.UniverseSettings.Leverage = 1
        #2. Required: Alpha Streams Models:
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
    
        #3. Required: Significant AUM Capacity
        self.SetCash(6000)
    
        #4. Required: Benchmark to SPY
        self.SetBenchmark("SPY")
        self.s1 = "AAPL"
        self.safe = "BIL"
        self.assets = [self.s1, self.safe]    # "TLT" as fix income in out-of-market period (high volatility)
        
        # Add Equity ------------------------------------------------ 
        for ticker in self.assets:
            self.AddEquity(ticker, Resolution.Daily).Symbol  #AddEquity
            
        # Set Scheduled Event Method For Kalman Filter updating.
        self.Schedule.On(self.DateRules.EveryDay(), 
            self.TimeRules.BeforeMarketClose(self.s1, 5), 
            self.EveryDayBeforeMarketClose)
            
            
    def EveryDayBeforeMarketClose(self):
        qb = self
        equ = self.AddEquity(self.s1, Resolution.Daily).Symbol
        # Get history
        history = qb.History(equ, 300, Resolution.Daily) #252*3
            
        # Get the close price daily return.
        close = history['close'].unstack(level=0)
        
        # Call pct_change to obtain the daily return
        returns = close.pct_change(7).iloc[7:]
        #self.sma_ratio = close.rolling(40).mean()/close.rolling(200).mean()

        
        # Initialize the HMM, then fit by the standard deviation data.
        model = MarkovRegression(returns, k_regimes=2, switching_variance=True).fit()
            
        # Obtain the market regime
        regime = model.smoothed_marginal_probabilities.values.argmax(axis=1)[-1]
    
        # ==============================
        
        if regime == 0:
            self.SetHoldings([PortfolioTarget(self.safe, 0.), PortfolioTarget(self.s1, 1.)])            
        else:
            self.SetHoldings([PortfolioTarget(self.safe, 0.), PortfolioTarget(self.s1, 0.)])


        #if regime == 0:
        #    self.SetHoldings([PortfolioTarget(self.safe, 0.), PortfolioTarget(self.s1, 1.)])            
        #else:
        #    self.SetHoldings([PortfolioTarget(self.safe, 1.), PortfolioTarget(self.s1, 0.)])