#
#   Compare Second/Minute resampled data vs Consolidated data 
#
#   Sets up a helper class SymbolData that allows different resample rates per symbol
#      self.assets holds a list of SymbolData objects to support multiple symbols 
#   Has optional delegates for OnDataConsolidated and OnSmaUpdated
#   Call pandas.resample when OnDataConsolidated fires
#   Uses a simple moving average indicator to Trade 
#      converted to use 5 minute sma bars 
#
#   Cleans up import * which is discouraged in python
#
#    You can view the QCAlgorithm base class on Github: 
#    https://github.com/QuantConnect/Lean/tree/master/Algorithm
#
# LICENSE is Apache 2.0: http://www.apache.org/licenses/LICENSE-2.0  
#
# https://github.com/QuantConnect/Lean/blob/dc27c4e3d655bd40da1f9d8b9b8ef850a2f10ee0/
#         Algorithm.Python/RollingWindowAlgorithm.py
#
# Modifications copyright (C) 2017 John Glossner 
# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# Is entire universe a public method of QCAlgorithm???
# https://www.quantconnect.com/lean/documentation/topic101.html
from QuantConnect import SecurityType, DataNormalizationMode, SymbolCache, Symbol, Chart 
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data import SubscriptionManager
from QuantConnect.Data.Market import TradeBar
from QuantConnect.Data.Consolidators import TradeBarConsolidator
from QuantConnect.Indicators import SimpleMovingAverage

import numpy as np
import pandas as pd 
from datetime import datetime, timedelta
from decimal import Decimal 

class ResampleAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.verbose = True     #For debugging

        # Backtest setup. Ignored in live trading 
        self.SetCash(100000)
        self.startDate = self.SetStartDate(2017,7,7)
        self.endDate = self.SetEndDate(2017,7,14) 
        #self.endDate = self.SetEndDate(2017,10,2)

        # Add assets with 5 minute consolidation  
        # self.symbols = [SymbolData(self, 'FB', resolution=Resolution.Minute, mode=DataNormalizationMode.Adjusted,
        #                           consolidated_window_length=10, consolidation_period=timedelta(minutes=5), 
        #                           sma_length=10)]

        # Can have multiple symbols with different parameters
        # FB=5min Consolidator SPY=10min Consolidator
        self.symbols = [SymbolData(self, 'FB', resolution=Resolution.Minute, mode=DataNormalizationMode.Adjusted,
                                  consolidated_window_length=10, consolidation_period=timedelta(minutes=5), 
                                  sma_length=10),
                        SymbolData(self, 'SPY', resolution=Resolution.Minute, mode=DataNormalizationMode.Adjusted,
                                  consolidated_window_length=15, consolidation_period=timedelta(minutes=10), 
                                  sma_length=20)]
                        
        #Schedule trading every 5 minutes
        self.Schedule.On(self.DateRules.EveryDay(), \
                 self.TimeRules.Every(timedelta(minutes=5)), \
                 Action(self.Trade))
                 
        return 
                        
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.'''
        pass

    def Resample(self, symbol_name, bars=3, frequency='5T'):
        '''
        Get historical bars and resample them to a higher frequency
        Uses pandas.resample
        
        Consolidators in QC use closed=left label=right
        QC 1min bars appear to use closed=left label=left (based on resampling 1sec data)
        To get QC 1m bars to match Quantopian 1min bars resample QC 1sec data with closed=right label=left 
        To get Quantopian data to match my broker's 5min data resample Q 1min data with closed=right label=left 
        To get QC 5min data to match my broker's 5min data resample QC 1min data with closed=left label=left 
        
        Note: The algorithm doesn't always work for open. 
              It seems to work for close. 
              I didn't check high/low/volumme 
        '''
        try:
            hist = self.History( [symbol_name], bars*5 )    #use list to get back a pandas datafram
                
            self.pretty_print_df(hist, -6)

            how = { 'open': 'first',
                    'high': 'max',
                    'low': 'min',
                    'close': 'last',
                    'volume': 'sum'}

            df = hist.resample(frequency, closed='right', label='left',level=1).agg(how).dropna()

            self.Log('after resample')
            self.pretty_print_df(df, -6)
            return         
            
        except Exception as e:
            self.Debug('@@@@@@@@@@@@@@@@@@@ EXCEPTION: Resample() @@@@@@@@@@@@@@@@@@@@@')
            self.Debug('Exception is %s' %e)
        
        #Invalid resample parameter
        return None
        
    def Trade(self):
        '''
        Determine if trade should be placed
        If current SMA > previous SMA: buy
        If current SMA < previous SMA: sell 
        '''
        for symbol in self.symbols:
            if not (symbol.consolidated_window_bars.IsReady and symbol.sma_window.IsReady): continue
            if not self.Securities[symbol.symbol_name].Exchange.ExchangeOpen: continue

            current_bar  = symbol.consolidated_window_bars[0]
            previous_bar = symbol.consolidated_window_bars[1]   #list is appended
            
            current_sma  = symbol.sma_window[0]
            previous_sma = symbol.sma_window[symbol.sma_window.Count-1]

            #Update plots
            self.Plot(symbol.symbol_name + " Chart", "SMA", current_sma.Value); 
            if self.Portfolio[symbol.symbol_name].IsLong:
                self.Plot(symbol.symbol_name + " Chart", "LongShort", Decimal(100.00));                
            if self.Portfolio[symbol.symbol_name].IsShort:
                self.Plot(symbol.symbol_name + " Chart", "LongShort", Decimal(-100.00));                

                            
            #Trade 
            if current_sma.Value > previous_sma.Value:
                if self.Portfolio[symbol.symbol_name].IsShort or not self.Portfolio[symbol.symbol_name].Invested:
                    self.SetHoldings(symbol.symbol_name, 1.0/len(self.symbols), liquidateExistingHoldings = True)
                    if self.verbose:
                        self.Log(' TRADE LONG ' + symbol.symbol_name + '\n' \
                            + '  prev bar time= ' + str(previous_bar.Time) \
                            + '  prev price= ' + '{0:.3f}'.format(previous_bar.Close) \
                            + '  curr bar time= ' + str(current_bar.Time) \
                            + '  curr price= ' + '{0:.3f}'.format(current_bar.Close) )
                        self.Log( '  prev sma time= ' + str(previous_sma.Time) \
                            + '  past sma= ' + '{0:.3f}'.format(previous_sma.Value) \
                            + '  curr sma time= ' + str(current_sma.Time) \
                            + '  curr sma= ' + '{0:.3f}'.format(current_sma.Value) )

            elif current_sma.Value < previous_sma.Value:
                if self.Portfolio[symbol.symbol_name].IsLong or not self.Portfolio[symbol.symbol_name].Invested: 
                    self.SetHoldings(symbol.symbol_name, -1.0/len(self.symbols), liquidateExistingHoldings = True)
                    if self.verbose:
                        self.Log(' TRADE SHORT ' + symbol.symbol_name + '\n' \
                            + '  prev bar time= ' + str(previous_bar.Time) \
                            + '  prev price= ' + '{0:.3f}'.format(previous_bar.Close) \
                            + '  curr bar time= ' + str(current_bar.Time) \
                            + '  curr price= ' + '{0:.3f}'.format(current_bar.Close) )
                        self.Log( '  prev sma time= ' + str(previous_sma.Time) \
                            + '  past sma= ' + '{0:.3f}'.format(previous_sma.Value) \
                            + '  curr sma time= ' + str(current_sma.Time) \
                            + '  curr sma= ' + '{0:.3f}'.format(current_sma.Value) )

    def pretty_print_df(self, df, length=5):
        '''
        Because what comes out of self.Log( str(df.tail())) looks like crap
        '''
        if not len(df): return 
    
        self.Log( '    '.join( df.columns.values ))
        if length > 0:
            if len(df) > length:
                count = 0 
                for index, row in df.iterrows():
                    if count > length: return
                    formatted_row = ["%.2f" % member for member in row]
                    self.Log( str(index) + '  ' + '  '.join(str(e) for e in formatted_row) )
                    count += 1 
            else:
                for index, row in df.iterrows():
                    formatted_row = ["%.2f" % member for member in row]
                    self.Log( str(index) + '  ' + '  '.join(str(e) for e in formatted_row) )
                
        if length < 0:
            count = abs(length)
            if len(df) > count:
                for index, row in df[::-1].iterrows():  #Reverse view into df 
                    if count ==0: return
                    formatted_row = ["%.2f" % member for member in row]
                    self.Log( str(index) + '  ' + '  '.join(str(e) for e in formatted_row) )
                    count -= 1 
            else:
                for index, row in df[::-1].iterrows():  #Reverse view into df 
                    if count ==0: return
                    formatted_row = ["%.2f" % member for member in row]
                    self.Log( str(index) + '  ' + '  '.join(str(e) for e in formatted_row) )

                        
class SymbolData(object):
    '''
    Packages Symbols to allow list like operations
    SymbolData should only be called from a class that inherits from QCAlgorithm 
    '''
    def __init__(self, qca, symbol_name, 
                 resolution = Resolution.Minute,
                 mode = DataNormalizationMode.Adjusted, #Raw, Adjusted, SplitAdjusted, TotalReturn
                 security_type=SecurityType.Equity, 
                 consolidated_window_length = None,
                 consolidation_period = timedelta(minutes=5),
                 sma_length = None ):
        self.qca = qca    #QCAlgorithm reference
        
        #Add equity to original qca
        self.symbol_name = symbol_name
        self.security_type = security_type
        self.symbol_object = qca.AddEquity(symbol_name, resolution) #QCAlgorithm public method 
        self.symbol_object.SetDataNormalizationMode(mode) 
        
        #Set up plots
        # Chart - Master Container for the Chart:
        stockPlot = Chart(symbol_name + " Chart")
        # On the Trade Plotter Chart we want 3 series: trades and price:
        stockPlot.AddSeries(Series("LongShort", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Price", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("SMA", SeriesType.Line, 0))
        qca.AddChart(stockPlot)
        
        
        #Set up consolidated tradebar window to keep historical data
        #Only setup if a window length given
        if consolidated_window_length:
            consolidator = TradeBarConsolidator(consolidation_period)
            consolidator.DataConsolidated += self.OnDataConsolidated
            qca.SubscriptionManager.AddConsolidator(symbol_name, consolidator)
            self.consolidation_period = consolidation_period
            self.consolidated_window_bars = RollingWindow[TradeBar](consolidated_window_length)

        #Creates an indicator and add to a rolling window when it is updated
        #Only setup if sma_length given
        if sma_length:
            #Why we need a symbol and not the string name is a mystery
            symbol = SymbolCache.GetSymbol(symbol_name)
            ind_name = qca.CreateIndicatorName(symbol, "SMA" + str(sma_length), resolution ) #inherited from QCAlgorithm
            
            self.sma = SimpleMovingAverage( ind_name, sma_length)
            sma_consolidator = TradeBarConsolidator(consolidation_period)   #Need a different delegate called
            qca.SubscriptionManager.AddConsolidator(symbol_name, sma_consolidator)
            qca.RegisterIndicator(symbol_name, self.sma, sma_consolidator)
            
            self.sma.Updated += self.OnSmaUpdated    #Called when sma updated 
            self.sma_window = RollingWindow[IndicatorDataPoint](sma_length) #store the last N sma values
        return 

    def OnDataConsolidated(self, sender, bar):
        '''
        delegate to update consolidate bars rolling window
        '''
        if self.qca.verbose:
            self.qca.Log(' CONSOLIDATOR ' +  self.symbol_name  + '  ' + str(self.qca.Time)  \
                    + ' open= ' +  '{0:.3f}'.format(bar.Open) \
                    + ' close= ' + '{0:.3f}'.format(bar.Close) \
                    + ' high= ' +  '{0:.3f}'.format(bar.High) \
                    + ' low= ' +   '{0:.3f}'.format(bar.Low) \
                    )
        # Add TradeBar in rollling window
        self.consolidated_window_bars.Add(bar)
        self.qca.Resample(self.symbol_name)
        return

    def OnSmaUpdated(self, sender, updated):
        '''
        Delegate to update SMA rolling window
        '''
        if self.qca.verbose:
            self.qca.Log(' SMA ' +  self.symbol_name  + '  ' + str(self.qca.Time)  \
                    + ' updated sma= ' +  str(updated) + '  type= ' + str( type(updated) ) \
                    + ' self.sma= ' + str(self.sma) + ' type self.sma= ' + str(type(self.sma)))
                    
        self.sma_window.Add(updated)
        return