# 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(2022, 1, 4)
        self.SetEndDate(2022, 1, 30)
        self.SetCash(1000000)
        
        self.underlying = self.AddIndex("SPX", Resolution.Hour)
        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.Hour)
        self.spxOptions.SetFilter(lambda u: (u.Strikes(0, 20).Expiration(0, 14)))

        # weekly option SPX contracts
        spxw = self.AddIndexOption(self.spx, "SPXW", Resolution.Hour)
        # set our strike/expiry filter for this option chain
        spxw.SetFilter(lambda u: (u.Strikes(0, 20)
                                     # 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)

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