# 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 *
from QuantConnect.Algorithm.CSharp import *

### <summary>
### Regression algorithm asserting we can specify a custom option exercise model
### </summary>
class CustomOptionExerciseModelRegressionAlgorithm(OptionAssignmentRegressionAlgorithm):
    def Initialize(self):
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        super().Initialize()

    def CustomSecurityInitializer(self, security):
        if Extensions.IsOption(security.Symbol.SecurityType):
            security.SetOptionExerciseModel(CustomExerciseModel())

    def OnData(self, data):
        super().OnData(data)

class CustomExerciseModel(DefaultExerciseModel):
    def OptionExercise(self, option: Option, order: OptionExerciseOrder):
        order_event = OrderEvent(
            order.Id,
            option.Symbol,
            Extensions.ConvertToUtc(option.LocalTime, option.Exchange.TimeZone),
            OrderStatus.Filled,
            Extensions.GetOrderDirection(order.Quantity),
            0.0,
            order.Quantity,
            OrderFee.Zero,
            "Tag"
        )
        order_event.IsAssignment = False
        return [ order_event ]

#region imports
from AlgorithmImports import *
#endregion
import pandas as pd
import random

from AlgorithmImports import *
from QuantConnect.Algorithm.CSharp import *

class CustomExerciseModel(DefaultExerciseModel):
    def OptionExercise(self, option: Option, order: OptionExerciseOrder):
        order_event = OrderEvent(
            order.Id,
            option.Symbol,
            Extensions.ConvertToUtc(option.LocalTime, option.Exchange.TimeZone),
            OrderStatus.Filled,
            Extensions.GetOrderDirection(order.Quantity),
            0.0,
            order.Quantity,
            OrderFee.Zero,
            "Tag"
        )
        order_event.IsAssignment = False
        return [ order_event ]


class VirtualYellowGiraffe(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2024, 4, 23)
        self.SetEndDate(2024, 4, 30)

        self.SetCash(100000) 
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.symbol = self.equity.Symbol
        
        self.InitOptionsAndGreeks(self.equity)
        ## Margin model
        #self.Portfolio.MarginCallModel = DefaultMarginCallModel(self.Portfolio, self.DefaultOrderProperties)
        self.Portfolio.MarginCallModel = MarginCallModel.Null
        self.SetSecurityInitializer(self.CustomSecurityInitializer)

        ## Options to target
        self.targetDELTA_STO = 0.03
        self.targetExpiryDTE_STO = 14
        self.minExpiryDTE = 7
        self.contracts = 10
        self.call_trade = False
        self.put_trade = False

    def CustomSecurityInitializer(self, security):
        if Extensions.IsOption(security.Symbol.SecurityType):
            security.SetOptionExerciseModel(CustomExerciseModel())

    ## Initialize Options settings, chain filters, pricing models, etc
    ## ====================================================================
    def InitOptionsAndGreeks(self, theEquity ):

        ## 1. Specify the data normalization mode (must be 'Raw' for options)
        theEquity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        ## 2. Set Warmup period of at leasr 10 days
        self.SetWarmup(15, Resolution.Daily)

        ## 3. Set the security initializer to call SetMarketPrice
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))

        ## 4. Subscribe to the option feed for the symbol
        theOptionSubscription = self.AddOption(theEquity.Symbol)

        ## 5. set the pricing model, to calculate Greeks and volatility
        theOptionSubscription.PriceModel = OptionPriceModels.CrankNicolsonFD()  # both European & American, automatically
                
        ## 6. Set the function to filter out strikes and expiry dates from the option chain
        theOptionSubscription.SetFilter(self.OptionsFilterFunction)

   

    def OnData(self, data):
        ## If we're done warming up, not invested
        if (not self.IsWarmingUp) and (not self.Portfolio.Invested) and (data.Bars.ContainsKey(self.symbol)):
            if (self.call_trade==False) and (self.put_trade==False) and (self.Time.hour >= 15):
                DTE_STO = self.targetExpiryDTE_STO
                ## Get contracts at delta=targetDELTA and DTE=targetExpiryDTE
                self.call_STO = self.SelectContract(self.equity.Symbol, self.targetDELTA_STO, 'delta', DTE_STO, OptionRight.Call)
                self.put_STO = self.SelectContract(self.equity.Symbol, self.targetDELTA_STO, 'delta', DTE_STO, OptionRight.Put)

                self.order(self.call_STO, self.contracts, 'STO')
                self.order(self.put_STO, self.contracts, 'STO')
                
                self.call_trade = True         
                self.put_trade = True
        elif (not self.IsWarmingUp) and (self.Portfolio.Invested) and (data.Bars.ContainsKey(self.symbol)):
            if (self.call_trade==True) and (self.put_trade==True):
                ## Liquidate if expiry of option is close
                if ((self.call_STO.Expiry - self.Time).days + 1 <= self.minExpiryDTE) and (self.Time.hour >= 14):
                    self.order(self.call_STO, self.contracts, 'BTC')
                    self.order(self.put_STO, self.contracts, 'BTC')
                    self.call_trade = False
                    self.put_trade = False

    ## Underlying symbol, delta, days till expiration, Option right (put or call)
    ## ============================================================================
    def order(self, symbolArg, contracts, ordertype= 'STO'):
        #self.current_delta =  self.Get_Current_Delta(symbolArg.Symbol)
        ## construct an order message -- good for debugging and order records
        if ordertype == 'STO':
            orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} |" + \
                            f"STO {symbolArg.Symbol} "
        elif ordertype == 'BTC':
            orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} |" + \
                            f"BTC {symbolArg.Symbol} "
                            
        #self.Debug(f"{self.Time} {orderMessage}")
        if ordertype == 'STO':
            self.Order(symbolArg.Symbol, -contracts, False, orderMessage)
        elif ordertype == 'BTC':
            self.Order(symbolArg.Symbol, contracts, False, orderMessage)

    ## Get an options contract that matches the specified criteria:
    ## Underlying symbol, delta, days till expiration, Option right (put or call)
    ## ============================================================================
    def SelectContract(self, symbolArg, strikeArg, strikeArgName, expiryDTE, optionRightArg= OptionRight.Call):

        canonicalSymbol = self.AddOption(symbolArg)
        theOptionChain  = self.CurrentSlice.OptionChains[canonicalSymbol.Symbol]
        theExpiryDate   = self.Time + timedelta(days=expiryDTE)
        
        ## Filter the Call/Put options contracts
        filteredContracts = [x for x in theOptionChain if x.Right == optionRightArg] 

        ## Sort the contracts according to their closeness to our desired expiry
        contractsSortedByExpiration = sorted(filteredContracts, key=lambda p: abs((p.Expiry+timedelta(hours=16)) - theExpiryDate), reverse=False)
        closestExpirationDate = contractsSortedByExpiration[0].Expiry
                                     
        ## Get all contracts for selected expiration
        contractsMatchingExpiryDTE = [contract for contract in contractsSortedByExpiration if contract.Expiry == closestExpirationDate]
        
        ## Get the contract with the contract with the closest delta
        if strikeArgName == 'delta':
            closestContract = min(contractsMatchingExpiryDTE, key=lambda x: abs(abs(x.Greeks.Delta)-strikeArg))

        return closestContract

    ## The options filter function.
    ## Filter the options chain so we only have relevant strikes & expiration dates. 
    ## =============================================================================
    def OptionsFilterFunction(self, optionsContractsChain):

        strikeCount  = 500 # no of strikes around underyling price => for universe selection
        min_Expiry_DTE = self.minExpiryDTE  # min num of days to expiration => for uni selection
        max_Expiry_DTE = self.targetExpiryDTE_STO  # max num of days to expiration => for uni selection
        
        ## Select options including both monthly and weekly
        return optionsContractsChain.IncludeWeeklys()\
                                    .Strikes(-strikeCount, strikeCount)\
                                    .Expiration(timedelta(min_Expiry_DTE), timedelta(max_Expiry_DTE))
   
   

#region imports
from AlgorithmImports import *
#endregion
class VirtualYellowGiraffe(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 4, 17)
        self.SetEndDate(2021, 2, 17)
        self.SetCash(100000) 
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        
        self.InitOptionsAndGreeks(self.equity)
        

    ## Initialize Options settings, chain filters, pricing models, etc
    ## ====================================================================
    def InitOptionsAndGreeks(self, theEquity ):

        ## 1. Specify the data normalization mode (must be 'Raw' for options)
        theEquity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        ## 2. Set Warmup period of at leasr 30 days
        self.SetWarmup(30, Resolution.Daily)

        ## 3. Set the security initializer to call SetMarketPrice
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))

        ## 4. Subscribe to the option feed for the symbol
        theOptionSubscription = self.AddOption(theEquity.Symbol)

        ## 5. set the pricing model, to calculate Greeks and volatility
        theOptionSubscription.PriceModel = OptionPriceModels.CrankNicolsonFD()  # both European & American, automatically
                
        ## 6. Set the function to filter out strikes and expiry dates from the option chain
        theOptionSubscription.SetFilter(self.OptionsFilterFunction)

    def OnData(self, data):
        
        ## If we're done warming up, and not invested, Sell a put. 
        if (not self.IsWarmingUp) and (not self.Portfolio.Invested):
            self.SellAnOTMPut()
      
    ## Sell an OTM Put Option.
    ## Use Delta to select a put contract to sell
    ## ==================================================================
    def SellAnOTMPut(self):
        
        ## Sell a 20 delta put expiring in 2 weeks (14 days)
        putContract = self.SelectContractByDelta(self.equity.Symbol, .30, 10, OptionRight.Put)
        
        ## construct an order message -- good for debugging and order rrecords
        orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} |" + \
                       f"Sell {putContract.Symbol} "+ \
                       f"({round(putContract.Greeks.Delta,2)} Delta)"
                       
        self.Debug(f"{self.Time} {orderMessage}")
        self.Order(putContract.Symbol, -1, False, orderMessage  )   
           
   
    ## Get an options contract that matches the specified criteria:
    ## Underlying symbol, delta, days till expiration, Option right (put or call)
    ## ============================================================================
    def SelectContractByDelta(self, symbolArg, strikeDeltaArg, expiryDTE, optionRightArg= OptionRight.Call):

        canonicalSymbol = self.AddOption(symbolArg)
        theOptionChain  = self.CurrentSlice.OptionChains[canonicalSymbol.Symbol]
        theExpiryDate   = self.Time + timedelta(days=expiryDTE)
        
        ## Filter the Call/Put options contracts
        filteredContracts = [x for x in theOptionChain if x.Right == optionRightArg] 

        ## Sort the contracts according to their closeness to our desired expiry
        contractsSortedByExpiration = sorted(filteredContracts, key=lambda p: abs(p.Expiry - theExpiryDate), reverse=False)
        closestExpirationDate = contractsSortedByExpiration[0].Expiry                                        
                                            
        ## Get all contracts for selected expiration
        contractsMatchingExpiryDTE = [contract for contract in contractsSortedByExpiration if contract.Expiry == closestExpirationDate]
    
        ## Get the contract with the contract with the closest delta
        closestContract = min(contractsMatchingExpiryDTE, key=lambda x: abs(abs(x.Greeks.Delta)-strikeDeltaArg))

        return closestContract

    ## The options filter function.
    ## Filter the options chain so we only have relevant strikes & expiration dates. 
    ## =============================================================================
    def OptionsFilterFunction(self, optionsContractsChain):

        strikeCount  = 100 # no of strikes around underyling price => for universe selection
        minExpiryDTE = 10  # min num of days to expiration => for uni selection
        maxExpiryDTE = 40  # max num of days to expiration => for uni selection
        
        return optionsContractsChain.IncludeWeeklys()\
                                    .Strikes(-strikeCount, strikeCount)\
                                    .Expiration(timedelta(minExpiryDTE), timedelta(maxExpiryDTE))