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

class MeanReversionDemo(QCAlgorithm):
    def initialize(self):
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2024, 9, 1)
        self.set_brokerage_model(BrokerageName.ALPHA_STREAMS)
        self.set_cash(1000000)
        self.set_benchmark("SPY")
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel())
        self.set_execution(ImmediateExecutionModel())

        # Add "SHY" to the assets list
        self.assets = ['PG', 'JNJ', 'MRK', 'GILD', 'SO', 'DUK', 'CL', 'BDX', 'BAESY', 'KMB', 
            'COR', 'GIS', 'HSY', 'ED', 'XEL', 'WEC', 'ORAN', 'CHT', 'AEE', 'CMS', 
            'CPB', 'CAG', 'SJM', 'TXNM', 'NEA', 'EQC','FLO', 'NAC', 'SAFT', 'EBF', 'SHY']

        for asset in self.assets:
            self.add_equity(asset, Resolution.MINUTE)

        # Schedule the event using "SHY"
        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, 120, 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()*0.9)).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:
            #avg = self.history(symbol, timedelta(days=126)).close.mean()
            #if avg >= self.securities[symbol].price:
            #    continue

            insights.append( Insight.price(symbol, timedelta(days=1), InsightDirection.UP, magnitude, confidence, None, weight) )
    
        self.emit_insights(insights)