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

### <summary>
### Basic template algorithm simply initializes the date range and cash. This is a skeleton
### framework you can use for designing an algorithm.
### </summary>
# Custom slippage implementation
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.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        self.SetStartDate(2017,1,1)
        self.SetEndDate(2017,2,5) 
        self.SetCash(30000)
        
        self.stocks = ["NUGT","DUST"]
        
        self.should_trade = False
        self.order_tickets = []
        self.open_trades = {}
        self.nearest_option = 27.0
        self.sell_with_days_left = 5.0
        self.layers = 4.0
        self.lev = 0.80
        
        self.short_call = True
        self.long_put = False
        self.short_put = False
        self.long_call = False
        
        self.weights = {}
        for stock in self.stocks:
            self.weights[stock] = 1.0/np.float(len(self.stocks))/self.layers
        
        for stock in self.stocks:
            self.AddEquity(stock, Resolution.Minute)
            self.Securities[stock].SetDataNormalizationMode(DataNormalizationMode.Raw)
            self.Securities[stock].SetSlippageModel(CustomSlippageModel())
            option = self.AddOption( stock, Resolution.Minute)
            # option.SetFilter(lambda u: u.IncludeWeeklys().Strikes(+3,+5).Expiration(timedelta(self.nearest_option), timedelta(30))) 
            option.SetFilter(self.UniverseFunc)
        
        
        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 UniverseFunc(self, universe):
        return universe.IncludeWeeklys().Strikes(+3, +5).Expiration(timedelta(self.nearest_option), timedelta(45))

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

    def get_stat( self, hist ):
        '''converts to log returns and get avereage return per time point, then sums'''
        for stock in self.stocks:
            if stock not in hist.index:
                return False
        hist1 = hist.loc[ self.stocks[0] ].close
        hist2 = hist.loc[ self.stocks[1] ].close
        
        new_guy =pd.DataFrame( { self.stocks[0]: hist1, self.stocks[1]: hist2 } )
        lr = np.diff( np.log( new_guy ), axis=0 )
        w = np.array( [ self.weights[ self.stocks[0] ], self.weights[ self.stocks[1] ] ] )
        avg_lr = np.sum( lr * w, axis=1 )
        stats = np.sum( avg_lr )
        
        return stats

    # Override the base class event handler for order events
    # def OnOrderEvent(self, orderEvent):
    #     order = self.Transactions.GetOrderById(orderEvent.OrderId)
    #     self.Debug("{0}: {1}: {2}".format(self.Time, order.Type, orderEvent))

    def OnData(self, data):
        if self.should_trade:
            nav = np.float(self.Portfolio.TotalPortfolioValue)
            trade_data = {}
            for stock in self.stocks:
                trade_data[stock] = {}
            
            #
            #calculate # of shares/dollars we want to trae based on leverage number
            for stock in self.stocks:
                price_stock = np.float( data[ stock ].Price )
                wanted_shares = -np.round( ( nav * self.lev * self.weights[ stock ] ) / price_stock )
                current_shares = np.float( self.Portfolio[ stock ].Quantity )
                trade_data[ stock ]['dollars_to_trade'] = ( wanted_shares - current_shares ) * price_stock
                trade_data[ stock ]['shares_to_trade'] = wanted_shares - current_shares
                trade_data[ stock ]['price'] = price_stock
            #
            #make stats about history of price movement of underlyings
            # stats = self.get_stat( self.History(self.stocks, 10, Resolution.Minute) )

            # if stats > 0.00:
            #     make_trade = True
            # else:
            #     make_trade = False

            # if make_trade == False:
            #     return
            make_trade = True
                
            #check current option holdings
            #do not hold options to expiration
            option_holding_underlyings = []
            options_to_close = []
            for sym in self.Portfolio.Keys:
                holding = self.Portfolio[sym]
                if holding.HoldStock and holding.Symbol.SecurityType == SecurityType.Option:
                    self.Log( "HOLDING: " + str(holding.Symbol) )
                    info = self.open_trades[ holding.Symbol ]
                    options_to_close.append([sym.Value, holding.Quantity])
            
            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
                
                #get underlying symbol and last price 
                
                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' ])
                
                # try to close trades here with market data
                for opt in options_to_close:
                    self.Log("got a trade to close")
                    self.Log( opt[0] )
                    
                    if opt[0] in df.index:
                        self.Log("looking at closing this")
                        p = df.loc[ opt[0] ]
                        qnty = opt[1]
                        self.Log( str( p ) )
                        
                        if int( ( p['expiry'] - data.Time ).days ) <= self.sell_with_days_left:
                            self.order_tickets.append( self.LimitOrder(p['Symbol'], -qnty, np.around( ( p['bid'] + p['ask'] ) / 2, 2) ) )
                        elif int( ( p['expiry'] - data.Time ).days ) > self.nearest_option - (self.nearest_option - self.sell_with_days_left)/self.layers:
                            if qnty > 0:
                                dir = "long"
                            else:
                                dir = "short"
                            
                            if p['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("----")
                        
                
                df = df.loc[ ( df['expiry'] > data.Time + timedelta(days=self.nearest_option) ) ]
                
                self.Log(usym+" "+str(upric))
                
                #get put
                puts = df.loc[ df['type'] == 1 ]
                
                # #get call
                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 ]['call2'] = calls.sort_values(['expiry','strike'], ascending=[False, False]).iloc[0]
                options[ usym ]['put'] = puts.sort_values(['expiry','strike'], ascending=[True, True]).iloc[0]
                options[ usym ]['qnty'] = np.round( self.lev * nav * self.weights[ usym ] / np.float( np.mean([options[ usym ]['call'].strike, options[ usym ]['put'].strike]) ) / 100 )
            
                
            #only trade if we have all the data
            
            for stock in self.stocks:
                if options[ stock ]['call'] is None or options[ stock ]['put'] is None or options[ stock ]['qnty'] is None:
                    return
            
            for stock in set(self.stocks):
                if [stock,'put','long'] not in option_holding_underlyings and self.long_put == True:
                    p = options[ stock ]['put']
                    qnty = options[ stock ]['qnty']
                if [stock,'put','short'] not in option_holding_underlyings and self.short_put == True:
                    p = options[ stock ]['put']
                    qnty = -options[ stock ]['qnty']
                if [stock,'call','short'] not in option_holding_underlyings and self.short_call == True:
                    p = options[ stock ]['call']
                    qnty = -options[ stock ]['qnty']
                # if [stock,'call','long'] not in option_holding_underlyings:
                #     c = options[ stock ]['call2']
                #     self.MarketOrder( c['Symbol'] , options[ stock ]['qnty'] )
                #     self.open_trades[ c['Symbol'] ] = c
                else:
                    continue
                self.Log("making that limir order")
                self.order_tickets.append( self.LimitOrder(p['Symbol'], qnty, np.around( p['bid'], 2) ) )
                self.open_trades[ p['Symbol'] ] = p

            self.should_trade = False