from Portfolio.BlackLittermanOptimizationPortfolioConstructionModel import BlackLittermanOptimizationPortfolioConstructionModel

from Execution.ImmediateExecutionModel import ImmediateExecutionModel

from Risk.NullRiskManagementModel import NullRiskManagementModel

from datetime import datetime, timedelta

class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        # Set requested data resolution
        self.UniverseSettings.Resolution = Resolution.Minute

        self.SetStartDate(2019, 1, 1)    #Set Start Date
        self.SetEndDate(2019, 1, 12)    #Set End Date
        self.SetCash(7750)               #Set Strategy Cash

        self.SetWarmUp(200)
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.ExtendedMarketHours = True
        tickers = ["TRVN", "AAPL", "BA", "TTP"]
        symbols = [ Symbol.Create(x, SecurityType.Equity, Market.USA) for x in tickers ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        
        self.SetAlpha(CustomEmaCrossAlphaModel(50, 200, Resolution.Minute))
        
        self.SetPortfolioConstruction(BlackLittermanOptimizationPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetRiskManagement(NullRiskManagementModel())
        

class CustomEmaCrossAlphaModel(AlphaModel):

    def __init__(self,
                 fastPeriod = 12,
                 slowPeriod = 26,
                 resolution = Resolution.Daily):
        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):

        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.Fast.IsReady and symbolData.Slow.IsReady:

                if symbolData.FastIsOverSlow:
                    if symbolData.Slow > symbolData.Fast:
                        insights.append(Insight.Price(symbolData.Symbol, timedelta(days=1), InsightDirection.Down))

                elif symbolData.SlowIsOverFast:
                    if symbolData.Fast > symbolData.Slow:
                        insights.append(Insight.Price(symbolData.Symbol, timedelta(days=1), InsightDirection.Up))

            symbolData.FastIsOverSlow = symbolData.Fast > symbolData.Slow

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):

        for added in changes.AddedSecurities:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                # create fast/slow EMAs
                symbolData = SymbolData(added, algorithm)
                symbolData.Fast = algorithm.EMA(added.Symbol, self.fastPeriod, self.resolution)
                symbolData.Slow = algorithm.EMA(added.Symbol, self.slowPeriod, self.resolution)
                symbolData.RegisterIndicator(algorithm)
                self.symbolDataBySymbol[added.Symbol] = symbolData
            else:
                # a security that was already initialized was re-added, reset the indicators
                symbolData.Fast.Reset()
                symbolData.Slow.Reset()


class SymbolData:

    def __init__(self, security, algorithm):
        self.Security = security
        self.Symbol = security.Symbol
        self.Fast = None
        self.Slow = None
        self.algorithm = algorithm
        ## Create consolidator for specific symbol here
        self.consolidator = TradeBarConsolidator(timedelta(days=1))    ## 5-period TradeBar Consolidator -- edit as
        self.consolidator.DataConsolidated += self.OnDataConsolidated  ## Add fuction to do what you want every 5-minutes with your data
        algorithm.SubscriptionManager.AddConsolidator(self.Symbol, self.consolidator)  ## Register consolidator

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        self.FastIsOverSlow = False

    def RegisterIndicator(self, algorithm):
        algorithm.RegisterIndicator(self.Symbol, self.Fast, self.consolidator)
        algorithm.RegisterIndicator(self.Symbol, self.Slow, self.consolidator)
        
    def OnDataConsolidated(self, sender, bar):
        
        ## Just example code to confirm the consolidator is working
        self.algorithm.Log("Consolidator Bar created for >> " + str(self.Symbol) + f"  Received at {self.algorithm.Time}")
        self.algorithm.Log("Fast EMA: " + str(round(self.Fast.Current.Value)))
        self.algorithm.Log("Slow EMA: " + str(round(self.Slow.Current.Value)))

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