#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.set_start_date(2014, 1, 1)
    
        #2. Required: Alpha Streams Models:
        self.set_brokerage_model(BrokerageName.ALPHA_STREAMS)
    
        #3. Required: Significant AUM Capacity
        self.set_cash(1000000)
    
        #4. Required: Benchmark to SPY
        self.set_benchmark("SPY")
        
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel())
        self.set_execution(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.add_equity(self.assets[i], Resolution.MINUTE).symbol
        
        # Set Scheduled Event Method For Our Model
        self.schedule.on(self.date_rules.every_day(), self.time_rules.before_market_close("SHY", 5), self.every_day_before_market_close)
        
        
    def every_day_before_market_close(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.emit_insights(insights)