# region imports
from AlgorithmImports import *
# endregion

class ATMStrangleAlgo(QCAlgorithm):
    def Initialize(self):
        # Algorithm settings
        self.SetStartDate(2025, 1, 1)
        self.SetEndDate(2025, 6, 30)
        self.SetCash(25000)
        self.set_brokerage_model(BrokerageName.CHARLES_SCHWAB, AccountType.MARGIN)
        self.default_order_properties.time_in_force = TimeInForce.DAY
        
        # Underlying equity and options
        self.symbol = self.AddEquity("QQQ", Resolution.Minute).Symbol
        option = self.AddOption("QQQ", Resolution.Minute)
        option.SetFilter(self.UniverseFunc)
        self.optionSymbol = option.Symbol

        # Risk parameters and thresholds
        self.limits = {
            'daily_range': 0.0052 # Take profit/close leg when range reached
        }
        self.targetDelta = 0.2  # Target delta for equal‑delta strangle
        self.stop_pct = -0.3    # stop out leg at -30%

        # Tracking variables
        self.currentPositions = {}    # {Symbol: info_dict}
        self.lastTradeDate = None
        self.dayOpenPrice = None      # Underlying open price for range checks

        # Schedule order entry and exit
        self.Schedule.On(self.DateRules.EveryDay(self.symbol),
                         self.TimeRules.AfterMarketOpen(self.symbol, 10),
                         self.OpenStrangle)

        self.Schedule.On(self.DateRules.EveryDay(self.symbol),
                         self.TimeRules.BeforeMarketClose(self.symbol, 5),
                         self.CloseAllPositions)

    def UniverseFunc(self, universe):
        # Include weeklies expiring in 1-7 days
        return universe.IncludeWeeklys().Expiration(timedelta(1), timedelta(7))

    def OpenStrangle(self):
        # Only trade once per day
        if self.lastTradeDate == self.Time.date(): return

        # Ensure no existing option positions
        if any(self.Portfolio[s].Invested for s in self.currentPositions): return

        # Get current option chain and underlying bar
        chain = self.CurrentSlice.OptionChains.get(self.optionSymbol)
        bar = self.CurrentSlice.Bars.get(self.symbol)
        if chain is None or bar is None: return

        underlying_price = bar.Close
        self.dayOpenPrice = underlying_price

        # Filter for valid contracts
        contracts = [c for c in chain if c.BidPrice>0 and c.AskPrice>0 and c.Greeks and c.Greeks.Delta is not None]
        if not contracts: return

        # Choose nearest expiration
        expiry = min({c.Expiry for c in contracts})
        daily_contracts = [c for c in contracts if c.Expiry == expiry]

        # Separate calls and puts
        calls = [c for c in daily_contracts if c.Right == OptionRight.Call]
        puts  = [c for c in daily_contracts if c.Right == OptionRight.Put]
        if not calls or not puts: return

        # Select legs by delta
        call_contract = self.SelectBestContract(calls, self.targetDelta, is_call=True)
        put_contract  = self.SelectBestContract(puts,  self.targetDelta, is_call=False)
        if not call_contract or not put_contract: return

        # Place 1-contract market orders
        self.MarketOrder(call_contract.Symbol, 1)
        self.MarketOrder(put_contract.Symbol,  1)

        # Record positions
        for contract, leg_type in [(call_contract, 'call'), (put_contract, 'put')]:
            mid = (contract.BidPrice + contract.AskPrice)/2
            self.currentPositions[contract.Symbol] = {
                'type': leg_type,
                'entry_price': mid
            }

        self.lastTradeDate = self.Time.date()
        self.Debug(f"[{self.time}] Opened ATM strangle at {underlying_price:.2f}: "
                   f"Call={call_contract.Symbol} (Δ={call_contract.Greeks.Delta:.2f}), "
                   f"Put={put_contract.Symbol} (Δ={put_contract.Greeks.Delta:.2f})")

    def SelectBestContract(self, contracts, target_delta, is_call):
        best = None; best_diff = float('inf')
        for c in contracts:
            d = c.Greeks.Delta
            diff = abs(d-target_delta) if is_call else abs(abs(d)-target_delta)
            if diff < best_diff:
                best_diff, best = diff, c
        return best

    def OnData(self, slice: Slice):
        # Only proceed if we have open positions and open price recorded
        if not self.currentPositions or self.dayOpenPrice is None:
            return

        # Underlying current price
        bar = slice.Bars.get(self.symbol)
        if bar is None: return
        current_price = bar.Close

        # Check for reaching average daily range
        move_pct = (current_price - self.dayOpenPrice) / self.dayOpenPrice
        if abs(move_pct) >= self.limits['daily_range']:
            # Directional close: if up, close call; if down, close put
            to_close = [s for s,info in self.currentPositions.items()
                        if (info['type']=='call' and move_pct>0) 
                        or (info['type']=='put' and move_pct<0)]
            for sym in to_close:
                if self.Portfolio[sym].Invested:
                    self.Liquidate(sym)
                    self.debug(f"Daily Range hit: closed {self.currentPositions[sym]['type']} leg {sym}")
                del self.currentPositions[sym]
            # After directional close, skip further stop checks this bar
            return

        # Check stop-loss of 50% on remaining legs
        for sym, info in list(self.currentPositions.items()):
            holding = self.Portfolio[sym]
            if not holding.Invested:
                del self.currentPositions[sym]
                continue
            pnl_pct = (holding.HoldingsValue - holding.HoldingsCost) / abs(holding.HoldingsCost)
            if pnl_pct <= self.stop_pct:
                self.Liquidate(sym)
                self.Debug(f"Stopped out {info['type']} leg {sym} at {pnl_pct:.1%} loss")
                del self.currentPositions[sym]

    def CloseAllPositions(self):
        total_pnl = 0
        for sym in list(self.currentPositions):
            if self.Portfolio[sym].Invested:
                # Get the exit price (midpoint of bid/ask if available, otherwise use last price)
                contract = self.Securities[sym]
                if contract.AskPrice > 0 and contract.BidPrice > 0:
                    exit_price = (contract.BidPrice + contract.AskPrice) / 2
                else:
                    exit_price = contract.Price

                entry_price = self.currentPositions[sym]['entry_price']
                quantity = self.Portfolio[sym].Quantity

                # Option contracts typically have a multiplier of 100
                leg_pnl = (exit_price - entry_price) * quantity * 100
                total_pnl += leg_pnl

                self.Debug(f"Closing {self.currentPositions[sym]['type']} {sym} | Entry: {entry_price:.2f}, Exit: {exit_price:.2f}, PnL: ${leg_pnl:.2f}")

                self.liquidate(sym)

        if total_pnl != 0:
            self.Log(f"End-of-day realized PnL: ${total_pnl:.2f}")

        self.currentPositions.clear()
        self.Log("Closed all positions at market close")


    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            order = self.Transactions.GetOrderById(orderEvent.OrderId)
            self.Debug(f"Order filled: {order.Symbol} {order.Quantity} @ {orderEvent.FillPrice:.2f}")