import numpy as np
class AlphaFiveUSTreasuries(QCAlgorithm):

    def Initialize(self):

        #1. Required: Five years of backtest history
        self.SetStartDate(2014, 1, 1)
    
        #2. Required: Alpha Streams Models:
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
    
        #3. Required: Significant AUM Capacity
        self.SetCash(1000000)

        #4. Required: Benchmark to SPY
        self.SetBenchmark("SPY")
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
    
        self.assets = ["IEF", "SHY", "TLT", "IEI", "SHV", "TLH", "EDV", "BIL",
                      "SPTL", "TBT", "TMF", "TMV", "TBF", "VGSH", "VGIT",
                      "VGLT", "SCHO", "SCHR", "SPTS", "GOVT"]

        self.symbols = {}
        
        # Add Equity ------------------------------------------------ 
        for i in range(len(self.assets)):
            self.symbols[self.assets[i]] = self.AddEquity(self.assets[i],Resolution.Minute).Symbol 
                
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen("IEF", 1), self.EveryDayAfterMarketOpen)

        
    def EveryDayAfterMarketOpen(self):
        qb = self
        # Fetch history on our universe
        df = qb.History(qb.Securities.Keys, 5, Resolution.Daily)
        # Make all of them into a single time index.
        df = df.close.unstack(level=0)
        # Calculate the truth value of the most recent price being less than 1 std away from the mean
        classifier = df.le(df.mean().subtract(df.std())).tail(1)
        # Get indexes of the True values
        classifier_indexes = np.where(classifier)[1]
        # Get the Symbols for the True values
        classifier = classifier.transpose().iloc[classifier_indexes].index.values
        # Get the std values for the True values (used for magnitude)
        magnitude = df.std().transpose()[classifier_indexes].values
        # Zip together to iterate over later
        selected = zip(classifier, magnitude)
        # ==============================
        
        insights = []
        for symbol, magnitude in selected:
            insights.append( Insight.Price(symbol, timedelta(days=5), InsightDirection.Up, magnitude) )
        self.EmitInsights(insights)
        
    def OnData(self, data):
        pass