from datetime import timedelta, date
import pandas as pd
### <summary>
### This example demonstrates how to add futures for a given underlying asset.
### It also shows how you can prefilter contracts easily based on expirations, and how you
### can inspect the futures chain to pick a specific contract to trade.
### </summary>
class BasicTemplateFuturesAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018,1,1)
        # self.SetEndDate(2016, 12, 23)
        self.SetCash(1000000)
        self.SetBenchmark('SPY')

        # Constants
        self.roll_days = 10  # See: https://coinmetrics.io/improving-your-futures-roll-an-overview-of-the-basis-trade/
        self.min_dur = 30 # Minimum duration of purchased future in days
        
        # Subscribe and set our expiry filter for the futures chain
        future = self.AddFuture(Futures.Currencies.BTC,Resolution.Minute)
        future.SetFilter(timedelta(0), timedelta(182))
        
        # Tracking
        self.fut_hold = None
        self.df = pd.DataFrame(columns=['fut','open','close'])
    
    def OnData(self,slice):
        if self.Time.date()==date(2020,12,23) and self.Time.hour==10 and self.Time.minute==0:
            self.Log(self.df)
        if self.Portfolio.Invested:
            if self.Time.hour==9 and self.Time.minute==30:
                contract = self.getfut(slice)
                self.df.loc[self.Time.date(),'fut'] = contract.Symbol 
                self.df.loc[self.Time.date(),'open'] = contract.LastPrice
            if self.Time.hour==16 and self.Time.minute==0:
                contract = self.getfut(slice)
                self.df.loc[self.Time.date(),'close'] = contract.LastPrice
            if (self.fut_hold.Expiry - self.Time).days <= self.roll_days:
                self.Liquidate()
        if not self.Portfolio.Invested:
            for chain in slice.FutureChains:
                # Get contracts expiring no earlier than in n days
                contracts = list(filter(lambda x: x.Expiry > self.Time + timedelta(self.min_dur), chain.Value))
                # Trade front contract
                if len(contracts) == 0: continue
                self.fut_hold = sorted(contracts, key = lambda x: x.Expiry, reverse=False)[0]
                self.SetHoldings(self.fut_hold.Symbol, 1)
                
    def getfut(self,slice):
        # get invested contract
        contract = self.fut_hold
        for chain in slice.FutureChains:
            filt = [i for i in chain.Value if i.Symbol==self.fut_hold.Symbol]
            if len(filt) == 0: continue
            contract = filt[0]
        
        return contract