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

class ModulatedMultidimensionalReplicator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)     
        #self.SetEndDate(2020, 1, 1)        
        self.SetCash(50000)

        tickers = ["QQQ","SPY","IYC","IYK","IGV","GLD","TLH","TLT"]
        
        for x in tickers:
            symbols = [ Symbol.Create(x, SecurityType.Equity, Market.USA) ]
            self.AddUniverseSelection( ManualUniverseSelectionModel(symbols) )
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddAlpha(MOMAlphaModel())
        
        self.Settings.RebalancePortfolioOnInsightChanges = False          
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(self.DateRules.Every(DayOfWeek.Monday)))
        self.SetExecution(ImmediateExecutionModel())    
        

        
    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 MOMAlphaModel(AlphaModel): 
    def __init__(self):
        
        self.indi = {} 
        self.indi_Filter = {}
        
        self.wind= 200
        self.num=3
        
        self.securities = [] 
        
    def OnSecuritiesChanged(self, algorithm, changes):
        
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            
            # initialize indicator
            self.indi[symbol] = CustomSlope( 'My_Custom', self.wind)
            algorithm.RegisterIndicator(symbol, self.indi[symbol], Resolution.Daily)
            
            # warmup indicator
            history = algorithm.History(symbol, self.wind, Resolution.Daily)    
            self.indi[symbol].WarmUp(history)
            
        # remove securities
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if security in self.indi:
                self.indi[symbol].remove(security)    
          
    def Update(self, algorithm, data):
        insights = []
        
        ## filtering
        filter1 = [x[0] for x in self.indi.items() if (self.indi[x[0]].Corr > 0.5) ]
        
        for symbol in self.indi:
            if symbol in filter1: # and filter2 and filter3
                self.indi_Filter[symbol] = self.indi[symbol]
        
        
        ## sorting
        ordered = sorted(self.indi_Filter.items(), key=lambda x: x[1].Value, reverse=True)[:self.num]        
        for x in ordered:
            symbol = x[0]
            insights.append( Insight.Price(symbol, timedelta(1), InsightDirection.Up) ) 
        
        # for testing
        algorithm.Plot("Custom_Slope", "Value", list(self.indi.values())[0].Value *10000)
        
        return insights
        
        

class CustomSlope:
    def __init__(self, name, period):
        self.Name = name
        self.Time = datetime.min
        self.IsReady = False
        self.Value = 0
        self.Slope = 0
        self.Corr = 0
        self.queue = deque(maxlen=period)
        
    def __repr__(self):
        return "{0} -> IsReady: {1}. Time: {2}. Value: {3}".format(self.Name, self.IsReady, self.Time, self.Value)
    
    def Update(self, input):
        #self.queue.appendleft(input.close)  # used by warm up
        self.queue.appendleft(input.Close)  # used by RegisterIndicator
        count = len(self.queue)
        #self.Time = input.Index # used by warm up
        self.Time = input.Time # used by RegisterIndicator
       
        self.IsReady = count == self.queue.maxlen
        
        #### start here the indicator calulation
        if self.IsReady:    
            y = np.log(self.queue)
            x = [range(len(y))]
            slope, corr = stats.linregress(x, y)[0], stats.linregress(x, y)[2]
            self.Slope = slope 
            self.Corr = corr  
            self.Value = slope * corr 
        #### finish the custom indicator
        
    def WarmUp(self,history):
        for tuple in history.itertuples():
            self.Update_warmup(tuple)

    def Update_warmup(self, input):
        self.queue.appendleft(input.close)  # used by warm up
        #self.queue.appendleft(input.Close)  # used by RegisterIndicator
        count = len(self.queue)
        self.Time = input.Index # used by warm up
        #self.Time = input.Time # used by RegisterIndicator
       
        self.IsReady = count == self.queue.maxlen
        
        #### start here the indicator calulation
        if self.IsReady:    
            y = np.log(self.queue)
            x = [range(len(y))]
            slope, corr = stats.linregress(x, y)[0], stats.linregress(x, y)[2]
            self.Slope = slope 
            self.Corr = corr  
            self.Value = slope * corr 
        #### finish the custom indicator