from datetime import timedelta
import numpy as np
from scipy import stats
from collections import deque

class ModulatedMultidimensionalReplicator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 6, 10)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        self.AddAlpha(MOMAlphaModel())
        
        self.AddEquity('SPY', Resolution.Daily)
        self.AddEquity('QQQ', Resolution.Daily)
        
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)
        

class CustomSlope:
    def __init__(self, name, period):
        self.Name = name
        self.Time = datetime.min
        self.IsReady = False
        self.Value = 0
        self.queue = deque(maxlen=period)
        
    # Update method is mandatory
    def Update(self, input):
        self.queue.append(input.Close)
        
        if not len(self.queue) == self.queue.maxlen:
            return False
        
        y= [np.log(data) for data in self.queue]
        x = [i for i in range(len(y))]
        slope,r_value = stats.linregress(x,y)[0],stats.linregress(x,y)[2]
        self.Value = slope
        
        return True
        
class MOMAlphaModel(AlphaModel): 
    def __init__(self):
        self.mom = {}       
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            
            self.mom[symbol] = CustomSlope('CSlope', 100)
            
            algorithm.RegisterIndicator(symbol, self.mom[symbol], Resolution.Daily) 
          
    def Update(self, algorithm, data):
        insights = []        
        ordered = sorted(self.mom.items(), key=lambda x: x[1].Value, reverse=True)[:4]        
        for x in ordered:
            symbol = x[0]
            insights.append( Insight.Price(symbol, timedelta(1), InsightDirection.Up) )    
        return insights