# region imports
from AlgorithmImports import *

from collections import deque
from statsmodels.tsa.stattools import acf
# endregion

class SmoothLightBrownDolphin(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 7, 10)  # Set Start Date
        self.SetStartDate(2020, 7, 20)  # Set Start Date
        self.SetCash(2000000)  # Set Strategy Cash
        
        self.sma = {}
        self.acf = {}
        # 1. Create the Autocorrelation indicator for each security
        self.amd = self.AddEquity("AMD", Resolution.Minute).Symbol
        self.sma[self.amd], self.acf[self.amd] = self.InitIndicators(self.amd)

        self.amzn = self.AddEquity("AMZN", Resolution.Minute).Symbol
        self.sma[self.amzn], self.acf[self.amzn] = self.InitIndicators(self.amzn)
        
        self.roku = self.AddEquity("ROKU", Resolution.Minute).Symbol
        self.sma[self.roku], self.acf[self.roku] = self.InitIndicators(self.roku)
        
        self.jpm = self.AddEquity("JPM", Resolution.Minute).Symbol
        self.sma[self.jpm], self.acf[self.jpm] = self.InitIndicators(self.jpm)

        self.CanTrade = set()
        self.MyInsights = []

        self.AddAlpha(CustomEmaCrossAlphaModel(self))

        self.SetWarmup(40)

    def InitIndicators(self,symbol):
        sma = SimpleMovingAverage('sma_'+symbol.Value,20)
        self.RegisterIndicator(symbol, sma, Resolution.Daily)
        
        acf_ind = CustomACF('ACF_'+symbol.Value,symbol,120,3)
        self.RegisterIndicator(symbol, acf_ind, Resolution.Daily)
        acf_ind.warmUpIndicator(self)
        return sma, acf_ind

    def OnData(self, data: Slice):
        if not self.Portfolio.Invested:
            for k,v in self.acf.items():
                if not v.IsReady:
                    return
            # # 2. One each indicator is ready get the stock with the maximun acf
            # symbol, acf_max = sorted([(k,v.MaxValue) for k,v in self.acf.items()],key= lambda x: x[-1])[-1]
            # # self.SetHoldings(symbol,0.5)
            # self.CanTrade.add(symbol)
            # self.Debug(f'The selected Maximun symbol was {symbol.Value}')
            
            # 3. One each indicator is ready get the stock with the minimum acf
            symbol, acf_min = sorted([(k,v.MinValue) for k,v in self.acf.items()],key= lambda x: x[-1])[0]
            self.CanTrade.add(symbol)
            # self.SetHoldings(symbol,0.5)
            self.Debug(f'The selected Minimum symbol was {symbol.Value}')
        if self.MyInsights:
            for insight in self.MyInsights:
                percent = 0.5 if insight.Direction == InsightDirection.Up else -0.5
                self.SetHoldings(insight.Symbol,percent)
        self.MyInsights = []
            
# 1. Create the Autocorrelation indicator for each security
class CustomACF(PythonIndicator):
    def __init__(self, name, symbol, period, nlags):
        self.Name = name
        self.Symbol = symbol
        self.WarmUpPeriod = period
        self.Time = datetime.min
        self.Value = 0
        self.Acf = None
        self.Price = deque(maxlen=period)
        self.LastTime = datetime.min + timedelta(minutes=1)
        self.nlags = nlags

    def warmUpIndicator(self, algorithm):
        history = algorithm.History(self.Symbol, self.WarmUpPeriod, Resolution.Daily).loc[self.Symbol]
        for idx, row in history.iterrows():
            self.Price.append(row.close) 

    @property
    def IsReady(self):
        return len(self.Price) == self.Price.maxlen

    @property
    def CurrentACF(self):
        if len(self.Price) == self.Price.maxlen:
            if self.Time == self.LastTime:
                return self.Acf
            self.LastTime = self.Time
            self.Acf = acf(self.Price, nlags=self.nlags)
            return self.Acf
        return None

    @property
    def MaxValue(self):
        if len(self.Price) == self.Price.maxlen:
            return max(self.CurrentACF)
    
    @property
    def MinValue(self):
        if len(self.Price) == self.Price.maxlen:
            return min(self.CurrentACF)
    
    @property
    def AverageValue(self):
        if len(self.Price) == self.Price.maxlen:
            return np.mean(self.CurrentACF)

    def Update(self, input: BaseData) -> bool:
        self.Price.append(input.Close)
        self.Time = input.Time
        return len(self.Price) == self.Price.maxlen

## ----------------- Modified SMA model
class CustomEmaCrossAlphaModel(AlphaModel):
    '''Alpha model that uses an EMA cross to create insights'''

    def __init__(self,
                 main_algo,
                 weight = 0.5,
                 fastPeriod = 20,
                 slowPeriod = 40,
                 resolution = Resolution.Daily):
        '''Initializes a new instance of the EmaCrossAlphaModel class
        Args:
            fastPeriod: The fast EMA period
            slowPeriod: The slow EMA period'''

        self.main_algo = main_algo
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        self.symbolDataBySymbol = {}

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{},{})'.format(self.__class__.__name__, fastPeriod, slowPeriod, resolutionString)


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.Fast.IsReady and symbolData.Slow.IsReady and symbol in self.main_algo.CanTrade:

                if symbolData.FastIsOverSlow:
                    if symbolData.Slow > symbolData.Fast:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down,weight=0.5))
                        self.main_algo.MyInsights.append(insights[-1])

                elif symbolData.SlowIsOverFast:
                    if symbolData.Fast > symbolData.Slow:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Up,weight=0.5))
                        self.main_algo.MyInsights.append(insights[-1])

            symbolData.FastIsOverSlow = symbolData.Fast > symbolData.Slow

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        for added in changes.AddedSecurities:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                symbolData = SymbolData(added, self.fastPeriod, self.slowPeriod, algorithm, self.resolution)
                self.symbolDataBySymbol[added.Symbol] = symbolData
            else:
                # a security that was already initialized was re-added, reset the indicators
                symbolData.Fast.Reset()
                symbolData.Slow.Reset()

        for removed in changes.RemovedSecurities:
            data = self.symbolDataBySymbol.pop(removed.Symbol, None)
            if data is not None:
                # clean up our consolidators
                data.RemoveConsolidators()

class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, security, fastPeriod, slowPeriod, algorithm, resolution):
        self.Security = security
        self.Symbol = security.Symbol
        self.algorithm = algorithm

        self.FastConsolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)
        self.SlowConsolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)

        algorithm.SubscriptionManager.AddConsolidator(security.Symbol, self.FastConsolidator)
        algorithm.SubscriptionManager.AddConsolidator(security.Symbol, self.SlowConsolidator)

        # create fast/slow EMAs
        self.Fast = SimpleMovingAverage(security.Symbol, fastPeriod)
        self.Slow = SimpleMovingAverage(security.Symbol, slowPeriod)

        algorithm.RegisterIndicator(security.Symbol, self.Fast, self.FastConsolidator);
        algorithm.RegisterIndicator(security.Symbol, self.Slow, self.SlowConsolidator);

        algorithm.WarmUpIndicator(security.Symbol, self.Fast, resolution);
        algorithm.WarmUpIndicator(security.Symbol, self.Slow, resolution);

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        self.FastIsOverSlow = False
        
    def RemoveConsolidators(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Security.Symbol, self.FastConsolidator)
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Security.Symbol, self.SlowConsolidator)

    @property
    def SlowIsOverFast(self):
        return not self.FastIsOverSlow