from AlgorithmImports import *

class WheelOptionStrategy(QCAlgorithm):
    '''
    Wheel Strategy on NVDA:
      - Sell 30-delta cash-secured puts ~30 days to expiration
      - Use on_order_event to log assignments, exercises, and expirations
      - Plot daily portfolio equity
      - Ensure no options are held on specific dates
    Start date: 2023-06-30
    '''

    def initialize(self):
        # Add underlying and option
        self._symbol = self.add_equity("NVDA", Resolution.HOUR).symbol
        self._option = self.add_option("NVDA", Resolution.HOUR)
        self._option.set_filter(self.universe_filter)

        # Backtest settings
        self.set_start_date(2023, 6, 30)
        self.set_cash(1000000)

        # Performance chart
        chart = Chart("Performance")
        chart.add_series(Series("Equity", SeriesType.LINE))
        self.add_chart(chart)

        # Strategy state
        self._put_order = None
        self._call_order = None
        self._has_stock = False

        # Schedule daily rolls ~30m after open
        self.schedule.on(
            self.date_rules.every_day(self._symbol),
            self.time_rules.after_market_open(self._symbol, 30),
            self.roll_positions
        )
        # Schedule option liquidation check daily before close
        self.schedule.on(
            self.date_rules.every_day(self._symbol),
            self.time_rules.before_market_close(self._symbol, 1),
            self.check_scheduled_liquidation
        )

    def universe_filter(self, universe):
        return universe.weeklys_only().expiration(timedelta(25), timedelta(35))

    def on_order_event(self, orderEvent):
        # Only process option orders
        symbol = orderEvent.Symbol
        if symbol.ID.SecurityType != SecurityType.OPTION:
            return

        msg = orderEvent.Message.lower() if orderEvent.Message else ""
        cash = self.portfolio.cash
        shares_held = self.portfolio[self._symbol].quantity

        # determine put or call
        right = symbol.ID.OptionRight
        kind = "put" if right == OptionRight.PUT else "call"

        if "assignment" in msg:
            self.debug(f"{self.time} - Assigned {kind} option: {symbol}. Shares held: {shares_held}. Cash: {cash}")
            self._has_stock = True if kind == "put" else False
            self._put_order = None
            self._call_order = None
        elif "exercise" in msg:
            self.debug(f"{self.time} - Exercised {kind} option: {symbol}. Shares held: {shares_held}. Cash: {cash}")
            self._has_stock = True if kind == "put" else False
            self._put_order = None
            self._call_order = None
        elif "otm" in msg:
            self.debug(f"{self.time} - {kind.capitalize()} option expired OTM: {symbol}. Shares held: {shares_held}. Cash: {cash}")
            self._put_order = None
            self._call_order = None

    def on_end_of_day(self, _symbol):
        if _symbol.ID.SecurityType != SecurityType.OPTION:
            return
        if self.time.date() != datetime(2024, 6, 7).date():
            return
        sec = _symbol
        if sec.Invested:
            self.debug(f"{self.time} - Holding {sec.Symbol}, Qty: {sec.Quantity}, Price: {sec.Price}")            

    def check_scheduled_liquidation(self):
        # Liquidate options on June 7, 2024
        if self.time.date() == datetime(2024, 6, 7).date():
            self.liquidate_options()

    def roll_positions(self):
        # Decide when to write new options
        if not self._put_order and not self._has_stock:
            self.sell_put()
        elif self._has_stock and not self._call_order:
            self.sell_call()

    def sell_put(self):
        chains = self.current_slice.option_chains.get(self._option.symbol)
        if not chains:
            return
        puts = [c for c in chains if c.right == OptionRight.PUT and c.expiry.date() > self.time.date()]
        if not puts:
            return
        contract = min(puts, key=lambda c: abs(c.greeks.delta + 0.30))
        max_contracts = int(self.portfolio.cash // (contract.strike * 100))
        if max_contracts <= 1:
            self.debug(f"{self.time} - Insufficient cash for put {contract.symbol}")
            return
        qty = -max_contracts
        self.market_order(contract.symbol, qty)
        self._put_order = (contract.symbol, contract.strike, max_contracts, contract.expiry.date())
        shares_held = self.portfolio[self._symbol].quantity
        cash = self.portfolio.cash
        self.debug(f"{self.time} - Sold put {contract.symbol} qty={qty}. Shares held: {shares_held}. Cash: {cash}")

    def sell_call(self):
        chains = self.current_slice.option_chains.get(self._option.symbol)
        if not chains:
            return
        calls = [c for c in chains if c.right == OptionRight.CALL and c.expiry.date() > self.time.date()]
        if not calls:
            return
        contract = min(calls, key=lambda c: abs(c.greeks.delta - 0.15))
        shares = self.portfolio[self._symbol].quantity
        contracts = shares // 100
        if contracts <= 0:
            self.debug(f"{self.time} - Skip call {contract.symbol}: no shares to cover")
            return
        qty = -contracts
        self.market_order(contract.symbol, qty)
        self._call_order = (contract.symbol, contract.strike, contracts, contract.expiry.date())
        shares_held = self.portfolio[self._symbol].quantity
        cash = self.portfolio.cash
        self.debug(f"{self.time} - Sold call {contract.symbol} qty={qty}. Shares held: {shares_held}. Cash: {cash}")

    def liquidate_options(self):
        # Force close all option positions
        for sym, holding in self.portfolio.items():
            if sym.SecurityType == SecurityType.OPTION and holding.Quantity != 0:
                qty = -holding.Quantity
                self.market_order(sym, qty)
                self.debug(f"{self.time} - Liquidated option {sym} qty={qty}. Cash: {self.portfolio.cash}")
        # reset state
        self._put_order = None
        self._call_order = None