# region imports
from AlgorithmImports import *
# endregion

from AlgorithmImports import *

class ChainedUniverseAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2023, 2, 2)
        # Need to set data normalization mode to raw for options to compare the strike price fairly.
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW

        self._contracts = {}
        # Filter for equities with fundamental data first.
        universe = self.add_universe(self._fundamental_function)
        # Based on the filtered equities, request an option universe with them as underlying.
        self.add_universe_options(universe, self._option_filter_function)

        symbol = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        self.schedule.on(
            self.date_rules.every_day(symbol),
            self.time_rules.every(timedelta(minutes=30)),
            self._select_option_contracts)

    def _select_option_contracts(self):
        def otm_filter(symbol, underlying_price):
            sid = symbol.id
            if sid.option_right == OptionRight.CALL and sid.strike_price > underlying_price:
                return True
            if sid.option_right == OptionRight.PUT and sid.strike_price < underlying_price:
                return True
            return False

        for underlying_symbol, contracts in self._contracts.items():
            price = self.securities[underlying_symbol].price
            to_add = [x for x in contracts if otm_filter(x, price)]
            for symbol in to_add:
                pass
                #self.add_option_contract(symbol)

    def _fundamental_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # Filter for equities with the lowest PE Ratio first.
        self._contracts.clear()
        filtered = [f for f in fundamental if not np.isnan(f.valuation_ratios.pe_ratio)]
        sorted_by_pe_ratio = sorted(filtered, key=lambda f: f.valuation_ratios.pe_ratio)
        return [Symbol.create("SPY", SecurityType.EQUITY, Market.USA)]  #[f.symbol for f in sorted_by_pe_ratio[:10]]

    def _option_filter_function(self, option_filter_universe: OptionFilterUniverse) -> OptionFilterUniverse:
        
        dte_min = 28
        dte_max = 91
        # We can filter by expiration here, since it doesn't change in a day
        option_filter_universe = option_filter_universe.include_weeklys().expiration(dte_min, dte_max)

        # Seed the underlying with previous day data 
        u = option_filter_universe.underlying
        security = self.securities[u.symbol]
        if not security.has_data:
            security.set_market_price(TradeBar(self.time, u.symbol, u.open, u.high, u.low, u.close, u.volume))
        
        # Use underlying symbol as dict key
        self._contracts[security.symbol] = [x.symbol for x in option_filter_universe]

        # Return empty filter
        return option_filter_universe.contracts([])
        
    def on_data(self, data: Slice) -> None:
        # Iterate each option chain to assert the contracts being selected.
        for symbol, chain in data.option_chains.items():
            for contract in chain:
                self.debug(f"Found {contract.symbol} option contract for {symbol}")