"""
https://www.quantconnect.com/forum/discussion/3245/using-option-greeks-to-select-option-contracts-to-trade

if you need Greeks:
 A) Filter and B) AddOption 
    more efficient than 
 C) OptionChainProvider and D) AddOptionContract
"""

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

class DeltaHedgedStraddleAlgo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2017, 12, 31)
        self.SetCash(1000000)
        
        self.Log("PERIOD: 2017")
        
        # ----------------------------------------------------------------------
        # Algo params
        # ----------------------------------------------------------------------
        self.PREMIUM = 0.01  # percentage of SPY share price
        self.MAX_EXPIRY = 30 # max num of days to expiration => for uni selection
        self._no_K = 2       # no of strikes around ATM => for uni selection
        self.resol = Resolution.Minute   # Resolution.Minute .Hour  .Daily
        self.tkr = "SPY"  # "SPY", "GOOG", ...
        self.Lev = d.Decimal(1.0)
        
        # 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_treshold = 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
        
        # Add options
        option = self.AddOption(self.tkr, self.resol)
        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(3))    # 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
        """
        # 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

        # 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:
            
            # select contract
            self.Log("get contract")
            self.get_contracts(slice)
            if (not self.call) or (not self.put): return
    
            # trade
            unit_price =  self.Securities[self.equity_symbol].Price * d.Decimal(100.0)   # share price x 100
            qnty = int(self.Portfolio.TotalPortfolioValue / unit_price)   
            # call_exists, put_exists = self.call is not None, self.put is not None
            
            if self.call is not None: self.Sell(self.call.Symbol, qnty)  # self.MarketOrder(self.call.Symbol, -qnty)
            if self.put is not None:  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_treshold:
                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

            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
            
            # 1. get furthest expiry            
            contracts_by_T = sorted(chain, key = lambda x: x.Expiry, reverse = True)
            if not contracts_by_T: return
            self.expiry = contracts_by_T[0].Expiry.date() # furthest expiry 
            self.last_trading_day = last_trading_day(self.expiry)
            
            # get contracts with further expiry and sort them by strike
            slice_T = [i for i in chain 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()) )

            # 2a. get the ATM closest CALL to short
            calls = [i for i in sorted_contracts \
                     if i.Right == OptionRight.Call and i.Strike >= spot_price]
            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))

            # 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[-1] if puts else None


    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]))

       
    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys()\
                    .Strikes(-self._no_K, self._no_K)\
                    .Expiration(timedelta(1), timedelta(self.MAX_EXPIRY))
    
    # ----------------------------------------------------------------------
    # Other ancillary fncts
    # ----------------------------------------------------------------------   
    def OnOrderEvent(self, orderEvent):
    #   self.Log("Order Event -> {}" .format(orderEvent))
        pass

    def OnAssignmentOrderEvent(self, assignmentEvent):
        self.Log(str(assignmentEvent))
        self._assignedOption = True
    #   if self.isMarketOpen(self.equity_symbol):
    #       self.Liquidate(self.equity_symbol)
    
    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
    
    # ----------------------------------------------------------------------
        # 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