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(2020, 7, 1)
        #self.SetEndDate(2020, 7, 30)
        self.SetCash(25000)
        
        # ----------------------------------------------------------------------
        # Algo params
        # ----------------------------------------------------------------------
        self.MAX_EXPIRY = 10 # max num of days to expiration => will select expiries upto 2 weeks out
        self.MIN_EXPIRY = 5 # max num of days to expiration => will select expiries at least 1 week out
        self._no_K = 50 # number of strikes below/above ATM for universe => No need to change this
        self.tkr = "SPY"  # "SPY", "GOOG", ... 
        self.quantity = 50 # number of credit spreads to trade
        
        consolidatedBar = TradeBarConsolidator(timedelta(hours=1)) # options data is minute by default, so hourly bar, change timedelta param for different bar lengths
        
        # bollinger band params
        self.bbLookback = 20
        self.bbBand = 2
        
        # profit targets
        self.profitTargetA,self.profitFactorA = 0.0,0.5 # profit factor the fraction of open to close i.e 0.5 = 50%
        self.profitTargetB = -0.05
        self.hedgeRatio = 1 / 20 # set as fraction for hedges per spread; set zero for no hedging
        
        self.shortDelta, self.hedgeDelta = -0.2,-0.3 # target deltas for short leg and hedge
        # ----------------------------------------------------------------------

        # add underlying Equity 
        self.resol = Resolution.Minute   # Resolution.Minute .Hour  .Daily
        equity = self.AddEquity(self.tkr, self.resol)  
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw) # IMPORTANT: default
        
        self.equity_symbol = equity.Symbol
        
        # Consolidate to 1hr bars and create bollingber band indicator
        self.SubscriptionManager.AddConsolidator(self.tkr, consolidatedBar)
        consolidatedBar.DataConsolidated += self.OnconsolidatedBar
        
        self.bb = BollingerBands(self.bbLookback,2) 
        self.RegisterIndicator(self.tkr, self.bb, consolidatedBar)
        
        # 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))
        self.SetSecurityInitializer(lambda x: x.SetBuyingPowerModel(CustomBuyingPowerModel(self)))
        
        # 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
        
        # Set warm up with 30 trading days to warm up the underlying volatility model
        self.SetWarmUp(30, Resolution.Daily)
        
        # vars for option contract objects
        self.shortLeg,self.longLeg,self.hedgeLeg = None, None, None
            
    
    
    
    def OnconsolidatedBar(self,sender,bar):
        if self.IsWarmingUp: return
        #self.Debug(str(self.Time) + " " + str(bar))
        
        # 2. sell puts if current price below bollinger band
        if not self.Portfolio.Invested:
            if self.bb.LowerBand.Current.Value > self.Securities[self.equity_symbol].Close:
                
                # select contracts
                self.get_contracts(self.CurrentSlice)
                if (not self.shortLeg) or (not self.longLeg) or (not self.hedgeLeg): return
                
                # get quantities for trading
                credit = self.Securities[self.shortLeg.Symbol].Price - self.Securities[self.longLeg.Symbol].Price
                
                qnty = self.quantity
                hedgeQnty = qnty * self.hedgeRatio # hedge (optional)
                self.Log(f"Thresh exceeded: Price {self.Securities[self.equity_symbol].Close} < {self.bb.LowerBand.Current.Value:.2f}. Credit sell to open value ${credit:.2f}")
                
                self.MarketOrder(self.shortLeg.Symbol,-qnty)
                self.MarketOrder(self.longLeg.Symbol,qnty)
                if hedgeQnty > 0: self.MarketOrder(self.hedgeLeg.Symbol,hedgeQnty)
                
                # set profit target credit value
                self.profitTargetA = - credit * self.profitFactorA
                
                
        
    def OnData(self, slice):
        if self.Portfolio.Invested:
            if (not self.shortLeg) or (not self.longLeg) or (not self.hedgeLeg): return
            # 3. Check exit conditions
            credit = self.Securities[self.shortLeg.Symbol].Price - self.Securities[self.longLeg.Symbol].Price
            #1. unrealized profit $ = arbitrary number
            if credit < self.profitTargetA:
                self.Liquidate()
                self.Log(f"Profit Target A hit {credit:.2f} < {self.profitTargetA:.2f}")
                
            #2. unrealized profit $ = -0.05
            if credit < self.profitTargetB:
                self.Liquidate()
                self.Log(f"Profit Target B hit {credit:.2f} < {self.profitTargetB:.2f}")
                
            #3. 0.5 delta on the short leg
            shortDelta = self.get_delta(slice,self.shortLeg) if self.get_delta(slice,self.shortLeg) is not None else 0
            if shortDelta < - 0.5:
                self.Liquidate()
                self.Log(f"Stop loss hit {shortDelta:.2f} < - 0.5 ( {credit:.2f} )")
            
            #4. Today is expiration day
            if self.Time.date() == self.last_trading_day and self.Time.hour == 15:
                self.Liquidate()
                self.Log(f"Liquidating on expiration ( {credit:.2f} )")
                
    def get_contracts(self, slice):
        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))
            
            # 1. get the contract series expiring on Friday
            contracts_by_T = [ i for i in chain if i.Expiry.date().weekday() == 4 ]
            
            #contracts_by_T = sorted(chain, key = lambda x: x.Expiry, reverse = False)
            #contracts_by_T = list(filter(lambda x: x.Expiry >= self.Time + timedelta(days = 30), chain))
            if not contracts_by_T: return
            
            #self.Log(f"Chain {[i.Expiry.date() for i in contracts_by_T]}")
        
            self.expiry = contracts_by_T[-1].Expiry.date() # furthest
            self.last_trading_day = last_trading_day(self.expiry)
            
            # sort contracts by strike
            sorted_contracts = sorted(contracts_by_T, key = lambda x: x.Strike, reverse = False)

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

            # get all puts below 50 delta
            puts = [i for i in sorted_contracts \
                     if i.Right == OptionRight.Put and i.Strike <= spot_price]
            
            # get short leg - the first strike below delta 
            self.shortLeg = [ i for i in puts if i.Greeks.Delta >= self.shortDelta ][-1] if puts else None
            
            # get long leg - the strike $2 below short leg
            self.longLeg = [ i for i in puts if i.Strike <= self.shortLeg.Strike - 2 ][-1] if puts else None
            
            # get hedge leg - 
            self.hedgeLeg = [ i for i in puts if i.Greeks.Delta >= self.hedgeDelta ][-1] if puts else None
            
            self.Log(f"Deltas: short {self.shortLeg.Greeks.Delta}, long {self.longLeg.Greeks.Delta}")
            
    def get_delta(self, slice,leg):
        """
        Take a data slice and an option leg as args; return the delta of that leg in slice
        """

        if not leg: 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 == leg.Symbol, chain)
            if not traded_contracts: self.Log("No traded cointracts"); return
        
            deltas = [i.Greeks.Delta for i in traded_contracts]
            
            return deltas[0]
            
    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys()\
                    .Strikes(-self._no_K, 0)\
                    .Expiration(timedelta(self.MIN_EXPIRY), timedelta(self.MAX_EXPIRY))
    
    # ----------------------------------------------------------------------
    # Other ancillary fncts
    # ----------------------------------------------------------------------   
    def OnOrderEvent(self, orderEvent):
    #   self.Log("Order Event -> {}" .format(orderEvent))
        pass

# necessitated by incorrect default margining of credit spreads in QC
class CustomBuyingPowerModel(BuyingPowerModel):
    def __init__(self, algorithm):
        self.algorithm = algorithm

    def HasSufficientBuyingPowerForOrder(self, parameters):
        # custom behavior: this model will assume that there is always enough buying power
        hasSufficientBuyingPowerForOrderResult = HasSufficientBuyingPowerForOrderResult(True)
        #self.algorithm.Log(f"CustomBuyingPowerModel: {hasSufficientBuyingPowerForOrderResult.IsSufficient}")
        return hasSufficientBuyingPowerForOrderResult

# ------------------------------------------------------------------------------
# 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