import pandas as pd
import datetime

class LogicalFluorescentOrangeButterfly(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2018,5,1)
        self.SetCash(1000000) 
        self.spy=self.AddEquity("SPY", Resolution.Minute)
        self.futures_bar_resolution=15 # 15 ie consolidate 1 min bar to 15min
        self.consolidator_dict={} # this is to keep track of the consolidator and hence can remove effciently
        self.generate_futures_dict()
        self.futures_child_to_parent_mapping_dict={}
        self.last_futures_active_contract_checked=0 # will only check if a future is active once a day
        self.last_spy_checked=0 # similar to above, but it is to check SPY

        
        
        
        
    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
        '''
        self.OnData_for_futures(data)
        
        if self.Time.hour>=9:
            if data.ContainsKey(self.spy.Symbol):
                if self.Time.date()!=self.last_spy_checked:
                    self.last_spy_checked=self.Time.date()
                    self.Plot('Custom', 'es', str(10*data.Bars[self.spy.Symbol].Close))
                    self.Plot('Custom', '/es', str(self.futures_data_dict['/ES']['dataframe']['close'].iloc[-1]))

                
    ## helpers
    def OnData_for_futures(self,data):
        if self.Time.date()!=self.last_futures_active_contract_checked:
            self.last_futures_active_contract_checked=self.Time.date()
            for chain in data.FutureChains.Values:
                contracts = chain.Contracts
                sorted_by_oi_contracts = sorted(contracts.Values, key=lambda k: k.OpenInterest, reverse=True)
                popular_contracts = sorted_by_oi_contracts[:3]
                active_contract=sorted_by_oi_contracts[0]
                if active_contract.Symbol!=self.futures_data_dict[str(chain.Symbol)]['active_symbol']:
                    last_contract_price=0
                    self.futures_child_to_parent_mapping_dict[str(active_contract.Symbol)]=str(chain.Symbol)
                    self.futures_data_dict[str(chain.Symbol)]['last_active_symbol']=self.futures_data_dict[str(chain.Symbol)]['active_symbol']
                    self.futures_data_dict[str(chain.Symbol)]['active_symbol']=str(active_contract.Symbol)
                    self.Log('active contract change for {}, from {} to {}'.format(str(chain.Symbol), self.futures_data_dict[str(chain.Symbol)]['last_active_symbol'], self.futures_data_dict[str(chain.Symbol)]['active_symbol']))
                    self.Log('new contract expiry is {}'.format(active_contract.Expiry))
                    if self.futures_data_dict[str(chain.Symbol)]['last_active_symbol']==None:
                        pass
                    else:
                        for tmp_contract in popular_contracts[1:]:
                            if str(tmp_contract.Symbol)==self.futures_data_dict[str(chain.Symbol)]['last_active_symbol']:
                                last_contract_price=tmp_contract.LastPrice
                        if last_contract_price==0: # 2017/12/31, ES front month contract has strange (nil?) data for a day
                            last_contract_price=self.futures_data_dict[str(chain.Symbol)]['dataframe']['close'][-1]
                        self.futures_data_dict[str(chain.Symbol)]['new_old_ratio']=active_contract.LastPrice / last_contract_price
                        self.futures_data_dict[str(chain.Symbol)]['dataframe']=self.futures_data_dict[str(chain.Symbol)]['dataframe']*self.futures_data_dict[str(chain.Symbol)]['new_old_ratio']
                        self.Log('attempting to delete existing symbol {}'.format(self.futures_data_dict[str(chain.Symbol)]['last_active_symbol']))
                        self.SubscriptionManager.RemoveConsolidator(self.futures_data_dict[str(chain.Symbol)]['last_active_symbol'], self.consolidator_dict[self.futures_data_dict[str(chain.Symbol)]['last_active_symbol']])
                        del self.consolidator_dict[self.futures_data_dict[str(chain.Symbol)]['last_active_symbol']]
                    self.consolidator_dict[self.futures_data_dict[str(chain.Symbol)]['active_symbol']] = TradeBarConsolidator(timedelta(minutes=self.futures_bar_resolution))
                    self.consolidator_dict[self.futures_data_dict[str(chain.Symbol)]['active_symbol']].DataConsolidated += self.Futures_OneMinuteBarHandler                        
                    self.SubscriptionManager.AddConsolidator(active_contract.Symbol, self.consolidator_dict[self.futures_data_dict[str(chain.Symbol)]['active_symbol']])
                    self.futures_data_dict[str(chain.Symbol)]['last_rolled_date']=self.Time.date()



                    

    def Futures_OneMinuteBarHandler(self, sender, bar):
        contract_symbol = sender.Consolidated.Symbol
        parent_symbol=self.futures_child_to_parent_mapping_dict[str(contract_symbol)]
        self.futures_data_dict[parent_symbol]['dataframe']=pd.concat([self.futures_data_dict[parent_symbol]['dataframe'], pd.DataFrame({'open':bar.Open, 'high':bar.High, 'low':bar.Low, 'close':bar.Close}, index=[bar.EndTime])],axis=0)


        

    def generate_futures_dict(self):
        self.futures_data_dict={}
        self.futures_data_dict['/ES']={'code':Futures.Indices.SP500EMini}
        #self.futures_data_dict['/MID']={'code':Futures.Indices.SP400MidCapEmini}
        self.futures_data_dict['/YM']={'code':Futures.Indices.Dow30EMini}
        #self.futures_data_dict['/RTY']={'code':Futures.Indices.Russell2000EMini}
        self.futures_data_dict['/GC']={'code':Futures.Metals.Gold}
        self.futures_data_dict['/SI']={'code':Futures.Metals.Silver}
        #self.futures_data_dict['/PL']={'code':Futures.Metals.Platinum}
        #self.futures_data_dict['/PA']={'code':Futures.Metals.Palladium}
        #self.futures_data_dict['/HG']={'code':Futures.Metals.Copper}
        
        for future in self.futures_data_dict.keys():
            self.futures_data_dict[future]['active_symbol']=None
            self.futures_data_dict[future]['last_rolled_date']=self.Time.date()-datetime.timedelta(days=100)
            self.futures_data_dict[future]['dataframe']=pd.DataFrame([],columns=['open','high','low','close'])
            tmp = self.AddFuture(self.futures_data_dict[future]['code'], Resolution.Minute)
            tmp.SetFilter(lambda universe: universe.Expiration(timedelta(0), timedelta(182)))