# region imports
from AlgorithmImports import *
# endregion


class WheelStrategyAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2015, 1, 1)
        self.set_cash(1_000_000)
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        self._equity = self.add_equity("SPY", data_normalization_mode=DataNormalizationMode.Raw)
        self._otm_threshold = self._default_otm_threshold = 0.15
        self._expiry = 28
        self._last_options_missing_log = None
        
    def _get_target_contract(self, right, target_price):
        # Start with default threshold
        self._otm_threshold = self._default_otm_threshold
        contract_symbols = self.option_chain_provider.get_option_contract_list(self._equity.symbol, self.time)

        if not contract_symbols:
            # Log missing options data (but not more frequently than once per day)
            if (self._last_options_missing_log is None or 
                self._last_options_missing_log.date() != self.time.date()):
                self.debug(f"No options data available for {self._equity.symbol} on {self.time.date()}")
                self._last_options_missing_log = self.time
            return None
        
        expiry = min([s.id.date for s in contract_symbols if s.id.date.date() >= self.time.date() + timedelta(self._expiry)])
        
        while True:
            filtered_symbols = [
                s for s in contract_symbols 
                if (s.id.date == expiry and s.id.option_right == right and 
                    (s.id.strike_price <= target_price if right == OptionRight.PUT else s.id.strike_price >= target_price))
            ]
            
            if filtered_symbols:
                symbol = sorted(filtered_symbols, key=lambda s: s.id.strike_price, reverse=right == OptionRight.PUT)[0]
                self.add_option_contract(symbol)
                return symbol
            
            # If no symbols found, reduce OTM threshold by 1% and try again
            self._otm_threshold -= 0.01
            
            # Set a minimum threshold to prevent infinite loop
            if self._otm_threshold <= 0.05:
                raise Exception(f"No suitable options found even with reduced threshold for {right}")
            
            # Recalculate target price with new threshold
            target_price = self._equity.price * (1-self._otm_threshold if right == OptionRight.PUT else 1+self._otm_threshold)

    def on_data(self, data):
        if not self.portfolio.invested and self.is_market_open(self._equity.symbol):
            symbol = self._get_target_contract(OptionRight.PUT, self._equity.price * (1-self._otm_threshold))
            if symbol:
                self.debug(f"Placed PUT order with OTM threshold: {self._otm_threshold}")
                self.set_holdings(symbol, -0.2)
        elif [self._equity.symbol] == [symbol for symbol, holding in self.portfolio.items() if holding.invested]:
            symbol = self._get_target_contract(OptionRight.CALL, self._equity.price * (1+self._otm_threshold))
            if symbol:
                self.debug(f"Placed CALL order with OTM threshold: {self._otm_threshold}")
                self.market_order(symbol, -self._equity.holdings.quantity / 100)