# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from AlgorithmImports import *

### <summary>
### This example demonstrates how to add and trade SPX index weekly options
### </summary>
### <meta name="tag" content="using data" />
### <meta name="tag" content="options" />
### <meta name="tag" content="indexes" />
class BasicTemplateSPXWeeklyIndexOptionsAlgorithm(QCAlgorithm):
   def Initialize(self):
      # self.SetStartDate(2020, 1, 17)  #Set Start Date
      # self.SetEndDate(2020, 1, 18)  #Set End Date
      # self.SetCash(50000)  #Set Strategy Cash
      # equity = self.AddEquity("GOOG", Resolution.Minute) # Add the underlying stock: Google
      # option = self.AddOption("GOOG", Resolution.Minute) # Add the option corresponding to underlying stock
      # self.symbol = option.Symbol
      # option.SetFilter(-10, 10, timedelta(0), timedelta(days = 1))
        
      self.SetStartDate(2022, 9, 1)
      self.SetEndDate(2022, 9, 30)
      self.SetCash(1000000)
      
      self.underlying = self.AddIndex("SPX", Resolution.Minute)
      self.underlying.SetDataNormalizationMode(DataNormalizationMode.Raw)
      self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

      self.spx = self.underlying.Symbol

      # regular option SPX contracts
      self.spxOptions = self.AddIndexOption(self.spx, Resolution.Minute)
      self.spxOptions.SetFilter(lambda u: (u.Strikes(-1, -1).Expiration(0, 0)))

      # weekly option SPX contracts
      spxw = self.AddIndexOption(self.spx, "SPXW", Resolution.Minute)
      # set our strike/expiry filter for this option chain
      spxw.SetFilter(lambda u: (u.Strikes(-1, 1)
                                    # single week ahead since there are many SPXW contracts and we want to preserve performance
                                    .Expiration(0, 1)
                                    .IncludeWeeklys()))

      self.spxw_option = spxw.Symbol

      self.SetSecurityInitializer(self.securityInitializer)
      self.currentDate = None

    # Called every time a security (Option or Equity/Index) is initialized
   def securityInitializer(self, security):
      security.SetDataNormalizationMode(DataNormalizationMode.Raw)
      security.SetMarketPrice(self.GetLastKnownPrice(security))
      if security.Type in [SecurityType.Option, SecurityType.IndexOption]:
         security.SetFillModel(BetaFillModel(self))
         #security.SetFillModel(MidPriceFillModel(self))
         security.SetFeeModel(TastyWorksFeeModel())

   def OnData(self,slice):
      # if it's the same day then stop.
      if self.Time.date().strftime("%Y-%m-%d") == self.currentDate:
         return

      for i in slice.OptionChains:
         if i.Key != self.spxw_option: 
            self.Log(f"\n({self.Time.strftime('%A')}) The `{i.Key}` does not match SPXW so we skip.")
            continue
      
         optionchain = i.Value
         df = pd.DataFrame([[x.Expiry] for x in optionchain],
                     index=[x.Symbol.Value for x in optionchain],
                     columns=['expiry'])
                     
         if df.shape[0] > 0:
            self.Log(f"\n{df.to_string()}")
            # we found data so we consider this day complete
            self.currentDate = self.Time.date().strftime("%Y-%m-%d")
         else:
            self.Log(f"\n({self.Time.strftime('%A')}) No options on this date")
            # self.Quit()


      # if self.Portfolio.Invested: return

      # chain = slice.OptionChains.GetValue(self.spxw_option)
      # if chain is None:
      #    return

      # # we sort the contracts to find at the money (ATM) contract with closest expiration
      # contracts = sorted(sorted(sorted(chain, \
      #    key = lambda x: x.Expiry), \
      #    key = lambda x: abs((chain.Underlying.Price + 20) - x.Strike)), \
      #    key = lambda x: x.Right, reverse=True)
      
      # # select just calls
      # contracts = [x for x in contracts if x.Right == OptionRight.Call]

      # # if found, buy until it expires
      # if len(contracts) == 0: return
      # symbol = contracts[0].Symbol
      # self.MarketOrder(symbol, -1)

   def OnOrderEvent(self, orderEvent):
      self.Debug(str(orderEvent))

class TastyWorksFeeModel:
   def GetOrderFee(self, parameters):
      optionFee = min(10, parameters.Order.AbsoluteQuantity * 0.5)
      transactionFee = parameters.Order.AbsoluteQuantity * 0.14
      return OrderFee(CashAmount(optionFee + transactionFee, 'USD'))

# Custom Fill model based on Beta distribution:
#  - Orders are filled based on a Beta distribution  skewed towards the mid-price with Sigma = bidAskSpread/6 (-> 99% fills within the bid-ask spread)
class BetaFillModel(ImmediateFillModel):

   # Initialize Random Number generator with a fixed seed (for replicability)
   random = np.random.RandomState(1234)
   
   def __init__(self, context):
      self.context = context
      
   def MarketFill(self, asset, order):
      # Start the timer
      # self.context.executionTimer.start()
   
      # Get the random number generator
      random = BetaFillModel.random
      # Compute the Bid-Ask spread
      bidAskSpread = abs(asset.AskPrice - asset.BidPrice)
      # Compute the Mid-Price
      midPrice = 0.5*(asset.AskPrice + asset.BidPrice)
      # Call the parent method
      fill = super().MarketFill(asset, order)
      # Setting the parameters of the Beta distribution:
      # - The shape parameters (alpha and beta) are chosen such that the fill is "reasonably close" to the mid-price about 96% of the times
      # - How close -> The fill price is within 15% of half the bid-Ask spread
      if order.Direction == OrderDirection.Sell:
         # Beta distribution in the range [Bid-Price, Mid-Price], skewed towards the Mid-Price
         # - Fill price is within the range [Mid-Price - 0.15*bidAskSpread/2, Mid-Price] with about 96% probability
         offset = asset.BidPrice
         alpha = 20
         beta = 1
      else:
         # Beta distribution in the range [Mid-Price, Ask-Price], skewed towards the Mid-Price
         # - Fill price is within the range [Mid-Price, Mid-Price + 0.15*bidAskSpread/2] with about 96% probability
         offset = midPrice
         alpha = 1
         beta = 20
      # Range (width) of the Beta distribution
      range = bidAskSpread/2.0
      # Compute the new fillPrice (centered around the midPrice)
      fillPrice = round(offset + range * random.beta(alpha, beta), 2)
      # Update the FillPrice attribute
      fill.FillPrice = fillPrice
      # Stop the timer
      # self.context.executionTimer.stop()
      # Return the fill
      return fill