import numpy as np
from datetime import timedelta
import pandas as pd

### <summary>
### Basic template algorithm simply initializes the date range and cash. This is a skeleton
### framework you can use for designing an algorithm.
### </summary>
class CustomSlippageModel:
    def GetSlippageApproximation(self, asset, order):
        return np.float( asset.Price ) * 0.002
        
class BasicTemplateAlgorithm(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''

    def Initialize(self):
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        self.SetStartDate(2016,1,1)  #Set Start Date
        self.SetEndDate(2016,3,1)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        self.stocks = ["DUST","NUGT"]
        self.should_trade = False
        self.order_tickets = {}
        self.open_trades = {}
        self.nearest_option = 27
        self.sell_with_days_left = 5
        self.short_call = True
        
        for stock in self.stocks:
            self.AddEquity(stock, Resolution.Minute)
            # self.Securities[stock].SetDataNormalizationMode(DataNormalizationMode.Raw)
            self.Securities[stock].SetSlippageModel(CustomSlippageModel())
            self.order_tickets[ stock ] = []
            option = self.AddOption( stock, Resolution.Minute)
            option.SetFilter(lambda u: u.IncludeWeeklys().Strikes(-1, +1).Expiration(timedelta(self.nearest_option), timedelta(45))) 
        
        self.lev = 1.0
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose(self.stocks[0], 10), Action(self.activate))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose(self.stocks[0], 6), Action(self.deactivate))

    def activate(self):
        self.should_trade = True
        
    def deactivate(self):
        self.should_trade = False

    def OnData(self, data):
        '''
        
        '''
        if self.should_trade:
            nav = np.float(self.Portfolio.TotalPortfolioValue)
            
            
            '''check current option holdings, do not hold options to expiration'''
            option_holding_underlyings = []
            for sym in self.Portfolio.Keys:
                holding = self.Portfolio[sym]
                if holding.HoldStock and holding.Symbol.SecurityType == SecurityType.Option:
                    info = self.open_trades[ holding.Symbol ]
                    if int( ( info['expiry'] - data.Time ).days ) <= self.sell_with_days_left:
                        self.SetHoldings(sym, 0)
                    else:
                        if holding.Quantity > 0:
                            dir = "long"
                        else:
                            dir = "short"
                        
                        if info['type'] == 1:
                            option_holding_underlyings.append( [str(holding.Symbol.Underlying), "put", dir] )
                        else:
                            option_holding_underlyings.append( [str(holding.Symbol.Underlying), "call", dir] )
            
                    self.Log( str( holding.Symbol ) )
            
            options = {}
            for stock in self.stocks:
                options[stock] = { 'put': None, 'call': None, 'qnty': None }
            
            for i in data.OptionChains:
                optionchain = list( i.Value )
                
                if len(optionchain) == 0:
                    continue
                
                try:
                    usym = str( optionchain[0].UnderlyingSymbol )
                    upric = float( optionchain[0].UnderlyingLastPrice )
                except:
                    return
                
                # create a data frame to show the filtered contracts
                # type(call 0, put 1)
                df = pd.DataFrame([[ x.Right, np.float(x.Strike), x.Expiry, np.float(x.BidPrice), np.float(x.AskPrice), x.Symbol] for x in optionchain ],
                                  index=[ x.Symbol.Value for x in optionchain ],
                                  columns=[ 'type', 'strike', 'expiry', 'ask', 'bid', 'Symbol' ])
                
                df = df.loc[ ( df['expiry'] > data.Time + timedelta(days=self.nearest_option) ) ]
                
                self.Log(usym+" "+str(upric))
                
                puts = df.loc[ df['type'] == 1 ]
                
                calls = df.loc[ df['type'] == 0 ]
                
                if len( calls.index ) < 1 or len( puts.index ) < 1:
                    continue
                
                options[ usym ]['call'] = calls.sort_values(['expiry','strike'], ascending=[True, False]).iloc[0]
                options[ usym ]['put'] = puts.sort_values(['expiry','strike'], ascending=[False, True]).iloc[0]
                options[ usym ]['qnty'] = np.round( self.lev * nav * 0.5 / np.float( np.mean([options[ usym ]['call'].strike, options[ usym ]['put'].strike]) ) / 100 )
            
            for stock in self.stocks:
                if [stock,'call','short'] not in option_holding_underlyings and self.short_call == True:
                    c = options[ stock ]['call']
                    self.Log( "trying to sell %s call contracts" % options[ stock ]['qnty'] )
                    self.Log( "at strike price %s" % c['strike'] )
                    self.MarketOrder( c['Symbol'] , -options[ stock ]['qnty'] )
                    self.open_trades[ c['Symbol'] ] = c
            
            self.should_trade = False