# region imports
from AlgorithmImports import *
from QuantConnect.Securities.Option import OptionPriceModels
# endregion

class CalmVioletWhale(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 12, 25)  # Set Start Date
        self.SetEndDate(datetime.now())
        self.SetCash(100000)  # Set Strategy Cash
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.symbol = self.equity.Symbol
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.SetBenchmark(self.symbol)

        self.SPYoption = self.AddOption("SPY", Resolution.Minute)

        # Strikes in the filter need to be made dynamic, a % on current price
        self.SPYoption.SetFilter(-150, 150, timedelta(0), timedelta(65))
        self.SPYoption.PriceModel = OptionPriceModels.BjerksundStensland()
        self.targetdelta = 0.2

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose("SPY", 10), self.OpenCondor)

        underlyingPriceChart = Chart("Underlying Price")
        self.AddChart(underlyingPriceChart)
        underlyingPriceChart.AddSeries(Series("Underlying Price", SeriesType.Line))
        underlyingPriceChart.AddSeries(Series("Call Strike Chosen", SeriesType.Line))

    def OnData(self, data: Slice):
        pass
    
    def OpenCondor(self):

        def RoundBaseFive(x, base=5):
            return base * round(x/base)
        
        longPut = None
        shortPut = None
        shortCall = None
        longCall = None

        # Dynamically filter available options based on current underlying price
        filterLowerBound = int(self.Securities[self.symbol].Price * 0.7)
        filterUpperBound = int(self.Securities[self.symbol].Price * 1.3)
        self.SPYoption.SetFilter(filterLowerBound, filterUpperBound, timedelta(0), timedelta(65))

        for i in self.CurrentSlice.OptionChains:
            chain = [x for x in i.Value]

            # If no option chain, return
            if (self.CurrentSlice.OptionChains.Count == 0):
                return

            # Sort option chain by expiration, choose furthest date
            expiry = sorted(chain, key = lambda x: x.Expiry)[-1].Expiry
            self.Log(str("Chosen expiry") + str(expiry))

            # filter calls and puts
            call = [i for i in chain if i.Right == OptionRight.Call and i.Expiry == expiry]
            put = [i for i in chain if i.Right == OptionRight.Put and i.Expiry == expiry]

            # sort calls by delta
            call_sorted = sorted(call, key = lambda x: (x.Strike))

            # get strike of call contract with target delta
            call_delta = min(call_sorted, key = lambda x: abs(x.Greeks.Delta - self.targetdelta)).Strike
            self.Log(str("Call Strike Chosen: ") + str(call_delta))

            # round to nearest base 5 to get a more liquid strike
            call_strike = RoundBaseFive(call_delta)
            self.Log(str("Underlying Price") + str(self.Securities[self.symbol].Price))

            # set long call strike +10 above
            call_strike_long = call_strike + 10

            # select both calls
            shortCall = min(call_sorted, key = lambda x: abs(x.Strike - call_strike))
            longCall = min(call_sorted, key = lambda x: abs(x.Strike - call_strike_long))

            # -- Same for puts --

            # sort puts by delta
            put_sorted = sorted(put, key = lambda x: (x.Strike), reverse=True)

            # get strike of put contract with target delta
            put_delta = min(put_sorted, key = lambda x: abs(x.Greeks.Delta <= -abs(self.targetdelta))).Strike
            self.Log(str("Put Strike Chosen: ") + str(put_delta))

            # Round to nearest base 5 to get more liquid strike
            put_strike = RoundBaseFive(put_delta)
            self.Log(str("Put Strike Rounded: ") + str(put_strike))

            # Set long put strike -10 below
            put_strike_long = put_strike - 10

            # Select both puts
            shortPut = min(put_sorted, key = lambda x: abs(x.Strike - put_strike))
            longPut = min(put_sorted, key = lambda x: abs(x.Strike - put_strike_long))

            shortPutDelta = shortPut.Greeks.Delta
            self.Log(str("Short Put Delta: ") + str(shortPutDelta))

            shortCallDelta = shortCall.Greeks.Delta
            self.Log(str("Short Call Delta: ") + str(shortCallDelta))

            # Size the position based on premiums

            # Open the condor

            self.Buy(longCall.Symbol, 1)
            self.Sell(shortCall.Symbol, 1)
            self.Sell(shortPut.Symbol, 1)
            self.Buy(longPut.Symbol, 1)

            # order logic? Limit order, etc


            self.Plot("Underlying Price", "Underlying Price", self.Securities[self.symbol].Price)
            self.Plot("Call Strike Chosen", "Call Strike Chosen", call_delta)