"""
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
# from datetime import datetime
import datetime

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 = 20       # no of strikes around ATM => for uni selection
        self.resol = Resolution.Minute  # Resolution.Minute .Hour  .Daily
        self.tkr = "GOOG"  # "SPY", "GOOG", ...
        self.Lev = d.Decimal(1.0)
        # self.last_trading_day = EndData
        self.last_trading_day = datetime.date(2020,1,1)
        # 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) 
        self.profit_target = d.Decimal(0.1)
        self.enter_trade_allowed = True
        # ----------------------------------------------------------------------

        # 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(-5, +5, 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 check_n_days_to_expiry(self, last_trading_day, today_date):
        current_date = datetime.strptime(today_date, '%Y-%m-%d').date()
        last_trading_date = datetime.strptime(last_trading_day, '%Y-%m-%d').date()
        if (last_trading_date-current_date < 5):
            return True
        else:
            return False

    def close_options(self):
        """ Liquidate opts (with some value) and underlying
        """
        if self.total_premium == 0.0:
            # self.Log("Close_options: Total premium is 0")
            return
        if (self.Portfolio.Invested is False):
            return
        today_date = self.Time.date()
        # self.Log("Today date: " + str(today_date))
        # self.Log("Last trading date: " + str(last_trading_day))
        # exit_cond = self.check_n_days_to_expiry(self.last_trading_day, today_date)
        # check this is the last trading day
        if self.last_trading_day != today_date:
            # self.Log("Close_options:Not the last trading day")
            total_ask_premium = d.Decimal(0)
            # total_ask_premium = 0.0
            
            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.Log("Close_options:x.Key: " + str(x.Key))
                    # self.Securities.ContainsKey(opt)
                    # only liquidate valuable options, otherwise let them quietly expiry
                    # total_ask_premium = total_ask_premium + float(self.Securities[x.Key].AskPrice)
                    ask_premium = d.Decimal(self.qnty) * self.Securities[x.Key].AskPrice
                    # self.Log("Close_options:ask_premium: "  + str(ask_premium))
                    total_ask_premium = total_ask_premium + ask_premium
                    # self.Log("Close_options:total_ask_premium: "  + str(total_ask_premium))
                    
            # self.Log("Close_options:Total premium = " + str(self.total_premium))
            if (total_ask_premium) < self.profit_target * self.total_premium:
                self.Log("Close_options:Total premium = " + str(self.total_premium))
                self.Log("Close_options:total_ask_premium: "  + str(total_ask_premium))
                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.Log("Close_options: liquidating: " + str(x.Key) )
                        self.Liquidate(x.Key)
                        
                # self.Log("Close_options:Returning after liquidating " )
                self.total_premium = 0.0
                return
    
        else:
            self.enter_trade_allowed = True
            self.Log("Close_options: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

        today_date = self.Time.date()
        if today_date > self.last_trading_day:
            self.enter_trade_allowed = True
        # if self.Portfolio.Invested:
        #     self.Log("Invested: get contract")
        #     # self.get_contracts_at_strike(slice)
        #     if (not self.call) or (not self.put): return
            
        # 2. sell options, if none
        if not self.Portfolio.Invested and self.enter_trade_allowed:
            # select contract
            # self.Log("get contract")
            self.get_contracts(slice)
            # self.Log("got  the contracts contract")
            if (not self.call) or (not self.put): return
            # self.Log("getting greeks")
            # self.get_greeks(slice)
            self.Log("OnData: Selected Call at strike:" + str(self.call.Strike) + " Delta:"
                    + str(self.call.Greeks.Delta))
            self.Log("OnData:Selected Put at strike:" + str(self.put.Strike) + " Delta:"
                    + str(self.put.Greeks.Delta))
            # trade
            unit_price =  self.Securities[self.equity_symbol].Price * d.Decimal(100.0)   # share price x 100
            self.Log("Price:" + str(self.Securities[self.equity_symbol].Price))
            self.Log("unit_price:" + str(unit_price))
            self.Log("Portfolio Value:" + str(self.Portfolio.TotalPortfolioValue))
            qnty = int(self.Portfolio.TotalPortfolioValue / unit_price)   
            # call_exists, put_exists = self.call is not None, self.put is not None
            self.Log("Qnty:" + str(qnty))
            self.qnty = qnty
            self.Log("Call Ask Price:" + str(self.call.AskPrice))
            self.Log("Put Ask Price:" + str(self.put.AskPrice))
            call_premium = qnty * self.call.AskPrice
            put_premium = qnty * self.put.AskPrice
            self.total_premium = call_premium + put_premium
            self.Log("Total premium = " + str(self.total_premium))
            self.enter_trade_allowed = False
            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)
            
            
        

    def get_contracts(self, slice):
        """
        Get call and put
        """
        # self.Log("Entered get_contracts")
        for kvp in slice.OptionChains:
            # self.Log("getting the key")
            if kvp.Key != self.option_symbol: continue
            # self.Log("got the key")
            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()) )

            calls = [i for i in sorted_contracts \
                     if i.Right == OptionRight.Call and i.Strike >= spot_price]
            deltas = [i.Greeks.Delta for i in calls]
            strikes = [i.Strike for i in calls]
            # self.Log("Call Strike: {}" .format(strikes))
            # self.Log("Call Delta: {}" .format(deltas))
            
            puts = [i for i in sorted_contracts \
                     if i.Right == OptionRight.Put and i.Strike <= spot_price]
            deltas = [i.Greeks.Delta for i in puts]
            strikes = [i.Strike for i in puts]
            # self.Log("Puts Strike: {}" .format(strikes))
            # self.Log("Puts Delta: {}" .format(deltas))
            
            for i in calls:
                if i.Greeks.Delta < 0.4:
                    self.call = i
                    self.Log("Selected Call at strike:" + str(self.call.Strike) + " Delta:"
                    + str(self.call.Greeks.Delta))
                    break
                
            # puts = sorted(puts, key=puts.Strike, reverse=False)
            for i in reversed(puts):    
                if i.Greeks.Delta >- 0.4:
                    self.put = i
                    self.Log("Selected Put at strike:" + str(self.put.Strike) + " Delta:"
                    + str(self.put.Greeks.Delta))
                    break
                
            # self.call = calls[0] if calls else None
            # 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]
            strikes = [i.Strike for i in traded_contracts]
            self.Log("Strike: {}" .format(strikes))
            self.Log("Delta: {}" .format(deltas))
            
            # self.Delta=sum(deltas)
            # self.Log("Vega: " + str([i.Greeks.Vega for i in traded_contracts]))
            # self.Log("Gamma: " + str([i.Greeks.Gamma for i in traded_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