#region imports
from AlgorithmImports import *
#endregion
# 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 QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta
import decimal as d
import random
from pprint import pprint

class CoveredCallAlgorithm(QCAlgorithm):

    def Initialize(self):
        self._no_K = 100       # no of strikes around ATM => for uni selection
        self.MIN_EXPIRY = 1 # min num of days to expiration => for uni selection
        self.MAX_EXPIRY = 180 # max num of days to expiration => for uni selection
        self.MAX_DELTA = d.Decimal(0.5)
        self.MIN_PREMIUM = d.Decimal(0.3)
        self.ticker = "SQQQ"
        #self.benchmarkTicker = "SPY"
        self.SetStartDate(2023, 2, 1)
        self.SetEndDate(2023, 2, 2)
        self.SetCash(100000)
        
        self.resolution = Resolution.Minute
        self.call, self.put, self.takeProfitTicket = None, None, None
        
        equity = self.AddEquity(self.ticker, self.resolution)
        option = self.AddOption(self.ticker, self.resolution)
        self.symbol = option.Symbol
        
        # set our strike/expiry filter for this option chain
        option.SetFilter(self.UniverseFunc)

        # 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
        
    def OnData(self,slice):
        # only run at a certain time
        if self.Time != datetime.strptime("2023-02-01 15:15:00", "%Y-%m-%d %H:%M:%S"):
            return

        # look for attractive put options
        self.TradePutOption(slice)

    def TradePutOption(self,slice):
        for i in slice.OptionChains:
            if i.Key != self.symbol: continue
        
            chain = i.Value
            
            # filter the put options contracts
            puts = [x for x in chain if x.Right == OptionRight.Put and x.BidPrice > self.MIN_PREMIUM] 

            # find best theta            
            contracts = sorted(puts, key = lambda x: x.Greeks.Theta, reverse=False)

            if len(contracts) == 0:
                self.Debug("No put contracts available, exiting")
                continue  
            
            # show the top 4 theta values and last price of underlying
            for i in range(0, min(4, len(contracts))):
                self.Log(f"{self.Time} {i} {contracts[i]} {contracts[i].UnderlyingLastPrice=} {contracts[i].Greeks.Theta=} {contracts[i].BidPrice=}")

    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys()\
                        .Strikes(-self._no_K, self._no_K)\
                        .Expiration(timedelta(self.MIN_EXPIRY), timedelta(self.MAX_EXPIRY))