class QuantumVerticalInterceptor(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 1, 7)
        self.SetCash(10000)
        self.DELTA_TARGET=0.5
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        self.df_calls = None
        #set a risk limit of 1% of portfolio value
        self.investment_limit = self.Portfolio.TotalPortfolioValue * 0.01 
        # Add the option
        option = self.AddOption("SPY")
        self.optionSymbol = option.Symbol

        # Add the initial contract filter
        option.SetFilter(-5, +5, 5, 10)

        # Define the Option Price Model
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        #option.PriceModel = OptionPriceModels.BlackScholes()
        #option.PriceModel = OptionPriceModels.AdditiveEquiprobabilities()
        #option.PriceModel = OptionPriceModels.BaroneAdesiWhaley()
        #option.PriceModel = OptionPriceModels.BinomialCoxRossRubinstein()
        #option.PriceModel = OptionPriceModels.BinomialJarrowRudd()
        #option.PriceModel = OptionPriceModels.BinomialJoshi()
        #option.PriceModel = OptionPriceModels.BinomialLeisenReimer()
        #option.PriceModel = OptionPriceModels.BinomialTian()
        #option.PriceModel = OptionPriceModels.BinomialTrigeorgis()
        #option.PriceModel = OptionPriceModels.BjerksundStensland()
        #option.PriceModel = OptionPriceModels.Integral()

        # Set warm up with 30 trading days to warm up the underlying volatility model
        self.SetWarmUp(30, Resolution.Daily)


    def OnData(self,slice):
        self.Plot("Portfolio", "Margin Remaining", self.Portfolio.MarginRemaining)         # Remaining margin on the account
        self.Plot("Portfolio", "Margin Used", self.Portfolio.TotalMarginUsed)         # Sum of margin used across all securities

        if self.IsWarmingUp or not slice.OptionChains.ContainsKey(self.optionSymbol):
            return

        chain = slice.OptionChains[self.optionSymbol]
        #set float format so delta displays correctly
        pd.set_option('display.float_format', lambda x: '%.5f' % x)
        #put the relevant data into the dataframe
        df = pd.DataFrame([[x.Right,float(x.Strike),x.Expiry,float(x.BidPrice),float(x.AskPrice),x.Greeks.Delta,x.UnderlyingLastPrice] for x in chain],
        index=[x.Symbol.Value for x in chain],
        columns=['type', 'strike', 'expiry', 'ask price', 'bid price', 'delta','underlyinglast'])
        #ensure expiry column is in datetime format
        df['expiry'] = pd.to_datetime(df['expiry'])
        # sort by expiry, descending
        df.sort_values(by=['expiry'],ascending=False)
        # get the most future date
        furthest_date = df['expiry'].iloc[0]
        # keep only those rows which have that furthest date
        df = df[df.expiry == furthest_date]
        #split the dataframe into calls and puts (calls are 0, puts are 1)
        self.df_calls = df[df.type==0]
        #sort by delta
        self.df_calls.sort_values(by=['delta'],ascending=False)

        #select the closest two records to the DELTA TARGET
        #try:
        uppercall_ind = self.df_calls[self.df_calls.delta<self.DELTA_TARGET].delta.idxmax()
        lowercall_ind = self.df_calls[self.df_calls.delta>self.DELTA_TARGET].delta.idxmin()
        self.df_calls  = self.df_calls[self.df_calls.index.isin([lowercall_ind,uppercall_ind])]
        spread_value = self.df_calls.at[lowercall_ind,'bid price'] - self.df_calls.at[uppercall_ind,'ask price']
        max_risk = self.df_calls.at[uppercall_ind,'strike'] - self.df_calls.at[lowercall_ind,'strike']
        max_risk_contract = max_risk * 100
        max_investment = math.trunc(self.investment_limit / max_risk_contract)
        self.Sell(OptionStrategies.BearCallSpread(self.optionSymbol, self.df_calls.at[lowercall_ind,'strike'], self.df_calls.at[uppercall_ind,'strike'] , self.df_calls.at[uppercall_ind,'expiry']), max_investment)
        #except:
        #    return

from QuantConnect.Securities.Option import OptionPriceModels
from QuantConnect.Securities import GetMaximumOrderQuantityForTargetBuyingPowerParameters
from datetime import timedelta
import decimal as d
from my_calendar import last_trading_day
import pandas as pd

class DeltaHedgedStraddleAlgo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 6)
        self.SetEndDate(2018, 4, 6)
        self.SetCash(100000)
        #self.Benchmark("SPY")
        self.Log("PERIOD: 2017-2019")
        
        # ----------------------------------------------------------------------
        # Algo params
        # ----------------------------------------------------------------------
        self.MIN_EXPIRY = 17 #Shortest DTE to consider
        self.MAX_EXPIRY = 37 #Longest DTE to consider
        self.MIN_DELTA = .37 #lowest delta = lower risk = less returns
        self.MAX_DELTA = .49 #highest delta = higher risk = more returns
        self.k_filter_low = -25
        self.k_filter_high = 1
        self.resol = Resolution.Minute   # Resolution.Minute .Hour  .Daily
        self.tkr = "SPY"  # "SPY", "GOOG", ...
        self.Lev = d.Decimal(1.0)
        self.single_trade_allocation = .4 #max portfolio value to place one single trade
        self.Settings.FreePortfolioValuePercentage = 0.07 #default is .025
    
        # self.Ntnl_perc = d.Decimal( round( 1. / (2. * self.MAX_EXPIRY/7.), 2) ) #  notional percentage, e.g. 0.08
        self.select_flag, self.hedge_flag = False, False
        self.previous_delta, self.delta_threshold = d.Decimal(0.0), d.Decimal(0.05) 
        # ----------------------------------------------------------------------

        
        # add underlying Equity 
        equity = self.AddEquity(self.tkr, self.resol)  
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw) # IMPORTANT: default
        self.equity_symbol = equity.Symbol
        
        #self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        # Add options
        option = self.AddOption(self.tkr, self.resol)
        option.SetDataNormalizationMode(DataNormalizationMode.Raw) # IMPORTANT: default
        #option.SetLeverage(2.0)
        self.option_symbol = option.Symbol

        # set our strike/expiry filter for this option chain
        option.SetFilter(self.UniverseFunc) # option.SetFilter(-2, +2, timedelta(0), timedelta(30))

        # for greeks and pricer (needs some warmup) - https://github.com/QuantConnect/Lean/blob/21cd972e99f70f007ce689bdaeeafe3cb4ea9c77/Common/Securities/Option/OptionPriceModels.cs#L81
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()  # both European & American, automatically
        # this is needed for Greeks calcs
        self.SetWarmUp(TimeSpan.FromDays(7))    # timedelta(7)
        
        self._assignedOption = False
        self.call, self.put = None, None
        
        # -----------------------------------------------------------------------------
        # scheduled functions
        # -----------------------------------------------------------------------------

        #self.Schedule.On(self.DateRules.EveryDay(self.equity_symbol), 
        #                 self.TimeRules.BeforeMarketClose(self.equity_symbol, 10),      
        #                 Action(self.close_options))

   
    def close_options(self):
        """ Liquidate opts (with some value) and underlying
        """
        #return
        # check this is the last trading day
        #if self.last_trading_day != self.Time.date(): return

        #self.Log("On last trading day: liquidate options with value and underlying ")

        # liquidate options (if invested and in the money [otherwise their price is min of $0.01)
        #for x in self.Portfolio:    # symbol = x.Key; security = x.Value ## but also symbol = x.Value.Symbol 
        #    if x.Value.Invested:    #  self.Portfolio[opt].Invested, but no need for self.Securities.ContainsKey(opt)
                # only liquidate valuable options, otherwise let them quietly expiry
        #        if self.Securities[x.Key].AskPrice > 0.05: self.Liquidate(x.Key)

        # CHECK if this necessary (incorporated above)
        if self.Portfolio[self.equity_symbol].Invested: 
            self.Liquidate(self.equity.Symbol)
            
    def OnData(self, slice):
        if self.IsWarmingUp: return
        if not self.HourMinuteIs(10, 1): return ##trade once a day max
        
        
        # choose ITM contracts
        #contracts = [x for x in call if call.UnderlyingLastPrice - x.Strike > 0]
        
        # or choose ATM contracts
        #contracts = sorted(optionchain, key = lambda x: abs(optionchain.Underlying.Price - x.Strike))[0]
        
        # or choose OTM contracts
        #contracts = [x for x in call if call.UnderlyingLastPrice - x.Strike < 0]
        
        # sort the contracts by their expiration dates
        #contracts = sorted(contracts, key = lambda x:x.Expiry, reverse = True)
        
        # 1. deal with any early assignments
        #if self._assignedOption:
            # close everything
        #    for x in self.Portfolio:  
        #        if x.Value.Invested: self.Liquidate(x.Key)
        self._assignedOption = False
        #   self.call, self.put = None, None  # stop getting Greeks

        # 2. sell options, if none
        if not self.Portfolio.Invested or self.Portfolio.GetBuyingPower(self.tkr) > self.Portfolio.TotalPortfolioValue * self.single_trade_allocation:
            
            
            
            # select contract
            #self.Log("get contracts")
            self.get_contracts(slice)
            if not self.put: return
    
            # trade
            #self.Securities[self.tkr].Holdings.MarginRemaining
            unit_price = self.put.Strike * d.Decimal(100.0)
            #unit_price2 = self.put.UnderlyingLastPrice * d.Decimal(100.0)
            #unit_price =  self.Securities[self.equity_symbol].Price * d.Decimal(100.0)   # share price x 100
            qnty = int((self.Portfolio.GetBuyingPower(self.tkr) * self.single_trade_allocation) / unit_price) 
            #q2 = int((self.Portfolio.GetBuyingPower(self.tkr) * .5) / unit_price)
            #qnty = min(q1, q2) #never use more than 50% of remaining buy power
            #bp = str(self.Portfolio.GetBuyingPower(self.tkr))
            #tv = str(self.Portfolio.TotalPortfolioValue)
            #mr = str(self.Portfolio.MarginRemaining)
            #mu = str(self.Portfolio.TotalMarginUsed)
            #cash = str(self.Portfolio.Cash) #settled only
            #u_cash = str(self.Portfolio.UnsettledCash) #unsettled only
            
            #blah = str(self.CalculateOrderQuantity(self.tkr, .2))

            ##fail
            #maxBpParam = GetMaximumOrderQuantityForTargetBuyingPowerParameters(self.Portfolio, self.Securities[self.equity_symbol], d.Decimal(20), False)
            #maxq = int(self.Portfolio.GetMaximumOrderQuantityForTargetBuyingPower(maxBpParam).Quantity)
            #
            #self.Log("unit_price " + str(unit_price))
            #self.Log("qty " + str(qnty))
            # call_exists, put_exists = self.call is not None, self.put is not None

            if self.call is not None and self.Portfolio[self.tkr].Invested:
                ccnt = self.Portfolio[self.tkr].Invested.Quantity / 100 
                self.Debug("Selling call " + str(ccnt) + "@" + str(self.call.Strike))
                self.Sell(self.call.Symbol, ccnt)  # self.MarketOrder(self.call.Symbol, -qnty)
            
            if self.put is not None and qnty > 0:
                self.Debug("Selling puts " + str(qnty) + "@" + str(unit_price))
                order = self.MarketOrder(self.put.Symbol, -qnty)
                


        # 3. delta-hedged any existing option
        #if self.Portfolio.Invested and self.HourMinuteIs(10, 1):
        #    self.get_greeks(slice)
        #    if abs(self.previous_delta - self.Delta) > self.delta_threshold:
        #        self.Log("delta_hedging: self.call {}, self.put {}, self.Delta {}" .format(self.call, self.put, self.Delta))
        #        self.SetHoldings(self.equity_symbol, self.Delta)
        #        self.previous_delta = self.Delta


    def get_contracts(self, slice):
        """
        Get ATM call and put
        """
        
        for kvp in slice.OptionChains:
            if kvp.Key != self.option_symbol: continue
            if not self.HourMinuteIs(10, 1): continue
            optionchain = kvp.Value   # option contracts for each 'subscribed' symbol/key 
            
            calls = [x for x in optionchain if x.Right == 0]
            self.call = calls[0] if calls else None
        #   self.Log("delta call {}, self.call type {}" .format(self.call.Greeks.Delta, type(self.call)))
        #   self.Log("implied vol {} " .format(self.call.ImpliedVolatility))
            #self.Debug(str(calls))
            #self.Debug(str(self.call))
            
            #self.Log("min delt " + str(self.MIN_DELTA))
            #self.Log("max delt " + str(self.MAX_DELTA))
            puts = [x for x in optionchain if x.Right == 1 and abs(x.Greeks.Delta) >= self.MIN_DELTA and abs(x.Greeks.Delta) <= self.MAX_DELTA]
            #df = pd.DataFrame([[x.Right, float(x.Strike), x.Expiry, float(x.BidPrice), float(x.AskPrice), x.Greeks.Delta, x.Greeks.Theta] for x in optionchain],
            #       index=[x.Symbol.Value for x in optionchain],
            #       columns=['type(call 0, put 1)', 'strike', 'expiry', 'ask', 'bid', 'delta', 'theta'])
            #put_frame = pd.DataFrame([[float(x.Strike), x.Expiry, float(x.BidPrice), float(x.AskPrice), x.Greeks.Delta, x.Greeks.Theta] for x in puts],
            #       index=[x.Symbol.Value for x in puts],
            #       columns=[ 'strike', 'expiry', 'ask', 'bid', 'delta', 'theta'])       
            #self.Log("spy close " + str(self.Securities[self.tkr].Close))
            #self.Log(str(put_frame))
            
        #for kvp in slice.OptionChains:
        #    if kvp.Key != self.option_symbol: continue
        #    chain = kvp.Value   # option contracts for each 'subscribed' symbol/key 
            
        #    spot_price = chain.Underlying.Price
        #   self.Log("spot_price {}" .format(spot_price))

            # prefer to do in steps, rather than a nested sorted
            ##
            ##TEN / DEL  *AKA* PRE / (DTE * DEL)
            ##
            
            # 1. get shortest expiry            
            contracts_by_T = sorted(puts, key = lambda x: x.Expiry, reverse = True)
            if not contracts_by_T: return
            self.expiry = contracts_by_T[-1].Expiry.date() # shortest expiry 
            self.last_trading_day = last_trading_day(self.expiry)
            
            # get contracts with shortest expiry and sort them by strike
            slice_T = [i for i in puts if i.Expiry.date() == self.expiry]
            sorted_contracts = sorted(slice_T, key = lambda x: x.Strike, reverse = False)

        #   self.Log("Expiry used: {} and shortest {}" .format(self.expiry, contracts_by_T[-1].Expiry.date()) )

     
            
            # 2b. get the ATM closest put to short
            #puts = [i for i in sorted_contracts \
            #         if i.Right == OptionRight.Put and i.Strike <= spot_price]
            self.put = puts[0] if puts else None 
            #self.Log("found contract: " + str(self.put))


    def get_greeks(self, slice):
        """
        Get greeks for invested option: self.call and self.put
        """

        if (self.call is None) or (self.put is None): return
        
        for kvp in slice.OptionChains:
            if kvp.Key != self.option_symbol: continue
            chain = kvp.Value   # option contracts for each 'subscribed' symbol/key 
            traded_contracts = filter(lambda x: x.Symbol == self.call.Symbol or 
                                         x.Symbol == self.put.Symbol, chain)
            if not traded_contracts: self.Log("No traded cointracts"); return
        
            deltas = [i.Greeks.Delta for i in traded_contracts]
        #   self.Log("Delta: {}" .format(deltas))
            self.Delta=sum(deltas)
            # self.Log("Vega: " + str([i.Greeks.Vega for i in contracts]))
            # self.Log("Gamma: " + str([i.Greeks.Gamma for i in contracts]))

       #.IncludeWeeklys()
    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys()\
                    .Strikes(self.k_filter_low, self.k_filter_high)\
                    .Expiration(timedelta(self.MIN_EXPIRY), timedelta(self.MAX_EXPIRY))
    
    # ----------------------------------------------------------------------
    # Other ancillary fncts
    # ----------------------------------------------------------------------   
    def OnOrderEvent(self, orderEvent):
    #   self.Log("Order Event -> {}" .format(orderEvent))
        pass
    
    def TimeIs(self, day, hour, minute):
        return self.Time.day == day and self.Time.hour == hour and self.Time.minute == minute
    
    def HourMinuteIs(self, hour, minute):
        return self.Time.hour == hour and self.Time.minute == minute
    
    def CustomSecurityInitializer(self, security):
        if security.Type == SecurityType.Option:
            security.MarginModel = OptionMarginModel()
    
    class MyPCM(InsightWeightingPortfolioConstructionModel):
        def CreateTargets(self, algorithm, insights):
            targets = super().CreateTargets(algorithm, insights)
            return [PortfolioTarget(x.Symbol, x.Quantity*algorithm.Securities[x.Symbol].Leverage) for x in targets]
            
    class OptionMarginModel:
        
        def __init__(self):
            optionMarginRequirement = 1;
            nakedPositionMarginRequirement = 0.1
            nakedPositionMarginRequirementOTM = 0.2   
        def Initialize(self):
            pass         
        def GetLeverage(security):
            return 1   
        def SetLeverage(security, leverage):
                return
    # ----------------------------------------------------------------------
        # all_symbols = [ x.Value for x in self.Portfolio.Keys ]
        # all_invested = [x.Symbol.Value for x in self.Portfolio.Values if x.Invested ]
        # for kvp in self.Securities: symbol = kvp.Key; security = kvp.Value
        #
        # orders = self.Transactions.GetOrders(None)
        # for order in orders: self.Log("order symbol {}" .format(order.Symbol))
        #
        # volatility = self.Securities[self.equity_symbol].VolatilityModel.Volatility
        # self.Log("Volatility: {}" .format(volatility))
        # set our strike/expiry filter for this option chain

# ------------------------------------------------------------------------------
# Business days
# ------------------------------------------------------------------------------
from datetime import timedelta #, date
from pandas.tseries.holiday import (AbstractHolidayCalendar,    # inherit from this to create your calendar
                                    Holiday, nearest_workday,   # to custom some holidays
                                    #
                                    USMartinLutherKingJr,       # already defined holidays
                                    USPresidentsDay,            # "     "   "   "   "   "
                                    GoodFriday,
                                    USMemorialDay,              # "     "   "   "   "   "
                                    USLaborDay,
                                    USThanksgivingDay           # "     "   "   "   "   "
                                    )


class USTradingCalendar(AbstractHolidayCalendar):
    rules = [
      Holiday('NewYearsDay', month=1, day=1, observance=nearest_workday),
      USMartinLutherKingJr,
      USPresidentsDay,
      GoodFriday,
      USMemorialDay,
      Holiday('USIndependenceDay', month=7, day=4, observance=nearest_workday),
      USLaborDay,
      USThanksgivingDay,
      Holiday('Christmas', month=12, day=25, observance=nearest_workday)
    ]

# TODO: to be tested
def last_trading_day(expiry):
    # American options cease trading on the third Friday, at the close of business 
    # - Weekly options expire the same day as their last trading day, which will usually be a Friday (PM-settled), [or Mondays? & Wednesdays?]
    # 
    # SPX cash index options (and other cash index options) expire on the Saturday following the third Friday of the expiration month. 
    # However, the last trading day is the Thursday before that third Friday. Settlement price Friday morning opening (AM-settled).
    # http://www.daytradingbias.com/?p=84847
    
    dd = expiry     # option.ID.Date.date()
    
    # if expiry on a Saturday (standard options), then last trading day is 1d earlier 
    if dd.weekday() == 5:
        dd -= timedelta(days=1)   # dd -= 1 * BDay()
        
    # check that Friday is not an holiday (e.g. Good Friday) and loop back
    while USTradingCalendar().holidays(dd, dd).tolist():    # if list empty (dd is not an holiday) -> False
        dd -= timedelta(days=1) 
        
    return dd