from AlgorithmImports import *
import datetime

class ShortStrangleWithStockCoverAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 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.Hour).Symbol
        self.option_symbol = self.AddOption(self.symbol, Resolution.Hour).Symbol
        self.contracts = []
        self.cover_orders = {}
        self.closing_orders = {}
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        self.SetWarmUp(40)
        self.lev = 2
        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, Resolution.Hour) #, 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(-5, 5).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.025 * self.Securities[self.symbol].Price
        put_otm_level = 0.975 * 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.nbr_contract = max(1, int(self.lev * self.Portfolio.Cash / self.Securities[self.symbol].Price / 100))

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

        self.cover_orders[self.contracts[0].Symbol] = self.StopMarketOrder(self.symbol, 100*self.nbr_contract, otm_calls[0].ID.StrikePrice)
        self.cover_orders[self.contracts[1].Symbol] = self.StopMarketOrder(self.symbol, -100*self.nbr_contract, otm_puts[0].ID.StrikePrice)
        if self.Portfolio[self.symbol].Quantity!=0 and (self.Securities[self.symbol].Price<otm_calls[0].ID.StrikePrice*0.985
            or self.Securities[self.symbol].Price>otm_puts[0].ID.StrikePrice*1.015):
            self.closing_orders[self.contracts[0].Symbol] = self.set_holdings(self.symbol, 0) #, otm_calls[0].ID.StrikePrice*0.995)
            #self.closing_orders[self.contracts[1].Symbol] = self.StopMarketOrder(self.symbol, 100, otm_puts[0].ID.StrikePrice*1.005)
        if self.Portfolio[self.symbol].Quantity!=0 and (self.Securities[self.symbol].Price>otm_calls[0].ID.StrikePrice*1.015
            or self.Securities[self.symbol].Price<otm_puts[0].ID.StrikePrice*0.985):
            self.closing_orders[self.contracts[0].Symbol] = self.Liquidate()

        self.Log(f"Opened strangle with contracts: {self.contracts[0].Symbol} (Call) and {self.contracts[1].Symbol} (Put)")


    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 OnOrderEvent(self, orderEvent):
     self.Log(f"OrderEvent: {orderEvent}")
     if orderEvent.Status == OrderStatus.Filled:
            order = self.Transactions.GetOrderById(orderEvent.OrderId)
            if order.Type == OrderType.StopMarket:
                if order.Symbol in self.cover_orders.values():
                    if order.Quantity > 0:
                        self.MarketOrder(order.Symbol, 100)
                    else:
                        self.MarketOrder(order.Symbol, -100)
                    self.cover_orders = {k: v for k, v in self.cover_orders.items() if v.OrderId != orderEvent.OrderId}
                elif order.Symbol in self.closing_orders.values():
                    if order.Quantity > 0:
                        self.MarketOrder(order.Symbol, -100)
                    else:
                        self.MarketOrder(order.Symbol, 100)
                    self.closing_orders = {k: v for k, v in self.closing_orders.items() if v.OrderId != orderEvent.OrderId}

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