from Risk.NullRiskManagementModel import NullRiskManagementModel


class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

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

        self.SetStartDate(2018, 1, 15)   #Set Start Date
        self.SetEndDate(2018, 7, 15)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        
        #self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.Resolution = Resolution.Daily
        symbols = [ Symbol.Create("IBM", SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        self.SetAlpha(HistoricalReturnsAlphaModel(resolution=Resolution.Daily, lookback=20, consolidationPeriod=7))
        self.SetPortfolioConstruction(NullPortfolioConstructionModel()) 
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())
        

    

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            # self.Debug("Purchased Stock: {0}".format(orderEvent.Symbol))
            pass

class HistoricalReturnsAlphaModel(AlphaModel):
    '''Uses Historical returns to create insights.'''

    def __init__(self, *args, **kwargs):
        '''Initializes a new default instance of the HistoricalReturnsAlphaModel class.
        Args:
            lookback(int): Historical return lookback period
            resolution: The resolution of historical data'''
        self.lookback = kwargs['lookback'] if 'lookback' in kwargs else 1
        self.resolution = kwargs['resolution'] if 'resolution' in kwargs else Resolution.Daily
        self.consolidationPeriod = kwargs['consolidationPeriod'] if 'consolidationPeriod' in kwargs else 1
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(self.resolution), self.lookback)
        self.symbolDataBySymbol = {}

    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.CanEmit:

                direction = InsightDirection.Flat
                magnitude = symbolData.Return
                if magnitude > 0: direction = InsightDirection.Up
                if magnitude < 0: direction = InsightDirection.Down
                insights.append(Insight.Price(symbol, self.predictionInterval, direction, magnitude, None))
                algorithm.Log(magnitude)
        
        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'''

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

        # initialize data for added securities
        symbols = [ x.Symbol for x in changes.AddedSecurities ]
        history = algorithm.History(symbols, self.lookback*self.consolidationPeriod, self.resolution)
        if history.empty: return

        tickers = history.index.levels[0]
        for ticker in tickers:
            symbol = SymbolCache.GetSymbol(ticker)

            if symbol not in self.symbolDataBySymbol:
                symbolData = SymbolData(symbol, self.lookback, self.consolidationPeriod)
                self.symbolDataBySymbol[symbol] = symbolData
                symbolData.RegisterIndicators(algorithm, self.resolution)
                symbolData.WarmUpIndicators(history.loc[ticker], algorithm)
                algorithm.Log(symbolData.Return)


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, symbol, lookback, consolidationPeriod):
        self.Symbol = symbol
        self.ROC = RateOfChange('{}.ROC({})'.format(symbol, lookback), lookback)
        # e.g. use a weekly conslidator to produce a weekly ROC 
        self.Consolidator = TradeBarConsolidator(TimeSpan.FromDays(consolidationPeriod))
        #algorithm.SubscriptionManager.AddConsolidator(self.Symbol, self.Consolidator)
        self.previous = None
        #algorithm.Debug("created symbol data with ROC period: " + str(lookback) + " conslidationPeriod: " + str(consolidationPeriod))

    def RegisterIndicators(self, algorithm, resolution):
        #self.Consolidator = algorithm.ResolveConsolidator(self.Symbol, resolution)
        algorithm.RegisterIndicator(self.Symbol, self.ROC, self.Consolidator)

    def RemoveConsolidators(self, algorithm):
        if self.Consolidator is not None:
            algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.Consolidator)

    def WarmUpIndicators(self, history, algorithm):
        #for index, tuple in history.itertuples():
            #self.ROC.Update(tuple.Index, tuple.close)
        
        for index, tuple in history.iterrows():
            tradeBar = TradeBar()
            tradeBar.Close = tuple['close']
            tradeBar.Open = tuple['open']
            tradeBar.High = tuple['high']
            tradeBar.Low = tuple['low']
            tradeBar.Volume = tuple['volume']
            tradeBar.Time = index
            self.Consolidator.Update(tradeBar)
        algorithm.Log("ROC is ready after warmup: " + str(self.ROC.IsReady))

    @property
    def Return(self):
        return float(self.ROC.Current.Value)

    @property
    def CanEmit(self):
        if self.previous == self.ROC.Samples:
            return False

        self.previous = self.ROC.Samples
        return self.ROC.IsReady

    def __str__(self, **kwargs):
        return '{}: {:.2%}'.format(self.ROC.Name, (1 + self.Return)**252 - 1)

class ImmediateExecutionModel(ExecutionModel):
    '''Provides an implementation of IExecutionModel that immediately submits market orders to achieve the desired portfolio targets'''

    def Execute(self, algorithm, targets):
        '''Immediately submits orders for the specified portfolio targets.
        Args:
            algorithm: The algorithm instance
            targets: The portfolio targets to be ordered'''
        for target in targets:
            open_quantity = sum([x.Quantity for x in algorithm.Transactions.GetOpenOrders(target.Symbol)])
            existing = algorithm.Securities[target.Symbol].Holdings.Quantity + open_quantity
            quantity = target.Quantity - existing
            if quantity != 0:
                algorithm.MarketOrder(target.Symbol, quantity)