from AlgorithmImports import *
import datetime

class ShortStrangleWithStockCoverAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 1, 1)  # Set Start Date
        self.SetEndDate(2022, 12, 31)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        self.symbol = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.option_symbol = self.AddOption(self.symbol).Symbol
        self.contracts = []
        self.cover_orders = {}
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        self.SetWarmUp(40)
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        #self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday),  self.on_data)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(15, 45), self.CheckStopOrders)
        option = self.add_option(self.symbol) #, self.Time)
        #option_chain.set_filter(min_expiry=timedelta(days=30), max_expiry=timedelta(days=40))
        option.set_filter(self.universe_func)
        self.symbolo = option.symbol

    def universe_func(self, universe: OptionFilterUniverse) -> OptionFilterUniverse:
        return universe.include_weeklys().strikes(-15, 15).expiration(timedelta(30), timedelta(40))    

    def on_data(self, slice: Slice) -> None:
      if self.Time.hour == 10 and self.Time.minute == 0:
                
        option_chain = self.OptionChainProvider.GetOptionContractList(self.symbolo, self.Time)
        if not option_chain:
            self.Log("No option chain available.")
            return
        
    

        #target_expiry = self.Time + timedelta(days=30)
        #expiry_dates = sorted(list(set([x.ID.Date for x in option_chain])))
        #expiry = min([date for date in expiry_dates if date >= target_expiry], key=lambda x: abs((x - target_expiry).days), default=None)

        call_otm_level = 1.02 * self.Securities[self.symbol].Price
        put_otm_level = 0.98 * self.Securities[self.symbol].Price
        #expiry = sorted([x.ID.Date for x in option_chain], key=lambda x: abs((x - self.Time).days - 30))
        otm_calls = sorted([x for x in option_chain if x.ID.OptionRight == OptionRight.Call and x.ID.StrikePrice > call_otm_level])
        otm_puts = sorted([x for x in option_chain if x.ID.OptionRight == OptionRight.Put and x.ID.StrikePrice < put_otm_level])

        if not otm_calls or not otm_puts:
            return

        self.contracts = [self.AddOptionContract(otm_calls[0], Resolution.Daily), self.AddOptionContract(otm_puts[0], Resolution.Daily)]
        self.Sell(self.contracts[0].Symbol, 1)
        self.Sell(self.contracts[1].Symbol, 1)

        self.cover_orders[self.contracts[0].Symbol] = self.StopMarketOrder(self.symbol, 100, otm_calls[0].ID.StrikePrice)
        self.cover_orders[self.contracts[1].Symbol] = self.StopMarketOrder(self.symbol, -100, otm_puts[0].ID.StrikePrice)
        
        self.Log(f"Opened strangle with contracts: {self.contracts[0].Symbol} (Call) and {self.contracts[1].Symbol} (Put)")

    def CheckStopOrders(self):
        if not self.contracts:
            return

        for contract in self.contracts:
            option = self.Securities[contract.Symbol]
            underlying_price = self.Securities[self.symbol].Price

            if underlying_price < self.contracts[0].StrikePrice and underlying_price > self.contracts[1].StrikePrice:
                self.LiquidateStockCover(contract)
            

    def LiquidateStockCover(self, contract):
        stock_quantity = self.Portfolio[self.symbol].Quantity
        if (contract.Symbol in self.cover_orders) and stock_quantity!=0:
            self.Liquidate(self.symbol)
            #del self.cover_orders[contract.Symbol]
            #self.contracts.remove(contract)
            #self.Buy(contract.Symbol, 1)  # Close the short option position
            self.Log(f"Liquidated stock cover for contract: {contract.Symbol}")

    def OnEndOfAlgorithm(self):
        self.Log(f"Final Portfolio Value: {self.Portfolio.TotalPortfolioValue}")