#region imports
from AlgorithmImports import *
from scipy.stats import norm, zscore
#endregion

class MeanReversionDemo(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(InsightWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
    
        self.assets = ["SHY", "TLT", "IEI", "SHV", "TLH", "EDV", "BIL",
                        "SPTL", "TBT", "TMF", "TMV", "TBF", "VGSH", "VGIT",
                        "VGLT", "SCHO", "SCHR", "SPTS", "GOVT"]
        
        # Add Equity ------------------------------------------------ 
        for i in range(len(self.assets)):
            self.AddEquity(self.assets[i], Resolution.Minute)
        
        # Set Scheduled Event Method For Our Model
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose("SHY", 5), self.EveryDayBeforeMarketClose)
        
        
    def EveryDayBeforeMarketClose(self):
        qb = self
        # Fetch history on our universe
        df = qb.History(qb.Securities.Keys, 30, Resolution.Daily)
        if df.empty: return
    
        # 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())).iloc[-1]
        if not classifier.any(): return
    
        # Get the z-score for the True values, then compute the expected return and probability
        z_score = df.apply(zscore)[[classifier.index[i] for i in range(classifier.size) if classifier.iloc[i]]]
    
        magnitude = -z_score * df.std() / df
        confidence = (-z_score).apply(norm.cdf)
    
        # Get the latest values
        magnitude = magnitude.iloc[-1].fillna(0)
        confidence = confidence.iloc[-1].fillna(0)
    
        # Get the weights, then zip together to iterate over later
        weight = confidence - 1 / (magnitude + 1)
        weight = weight[weight > 0].fillna(0)
        sum_ = np.sum(weight)
        if sum_ > 0:
            weight = (weight) / sum_
            selected = zip(weight.index, magnitude, confidence, weight)
        else:
            return
    
        # ==============================
        
        insights = []
        
        for symbol, magnitude, confidence, weight in selected:
            insights.append( Insight.Price(symbol, timedelta(days=1), InsightDirection.Up, magnitude, confidence, None, weight) )
    
        self.EmitInsights(insights)