# region imports
from AlgorithmImports import *
from itertools import groupby
# endregion

class MuscularVioletFox(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2024, 1, 1)

        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        date_rules = self.date_rules.every(DayOfWeek.MONDAY)
        self.universe_settings.schedule.on(date_rules)
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self._universe = self.add_universe(self._fundamental_selection_function)

        # Schedule a filtering on Monday morning every week
        spy = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        self.schedule.on(date_rules, self.time_rules.after_market_open(spy, 30), self.filter_and_trade)

    def _fundamental_selection_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # Filter for securities with (price > $10), has fundamental data, and has a price_to_earnings ratio. 
        filtered = [f for f in fundamental if f.price > 10 and f.has_fundamental_data and not np.isnan(f.valuation_ratios.pe_ratio)]
        # Sort filtered securities by dollar volume in decending order, keep the top 100 (securities with the most dollar volume). 
        sorted_by_dollar_volume = sorted(filtered, key=lambda f: f.dollar_volume, reverse=True)[:100]
        # Sort the previously sorted 100 by P/E ratio, keep the securities with the lowest 10 P/E ratios. 
        sorted_by_pe_ratio = sorted(sorted_by_dollar_volume, key=lambda f: f.valuation_ratios.pe_ratio, reverse=False)[:10]
        # Return the final selected securities
        return [f.symbol for f in sorted_by_pe_ratio]

    def on_securities_changed(self, changes):
        for security in [x for x in changes.added_securities if x.type==SecurityType.EQUITY]:
            security.dividend_yield_provider = DividendYieldProvider(security.symbol)
            security.contracts = []
        # When we remove the under
        for security in [x for x in changes.removed_securities if x.type==SecurityType.EQUITY]:
            self.liquidate(security.symbol)
            for contract in security.contracts:
                self.remove_option_contract(contract)

    def filter_and_trade(self) -> None:
        keyfunc = lambda x: x.id.strike_price
        option_model = OptionPricingModelType.FORWARD_TREE

        for symbol in self._universe.selected:
            # Get all option contracts trading
            contract_symbols = self.option_chain_provider.get_option_contract_list(symbol, self.time)
            if not contract_symbols:
                message = f'No options for {symbol} on {self.time}'
                #self.log(message)
                continue
            # Get the first expiry that just over 7 days later and within 1 year
            expiry = min(x.id.date for x in contract_symbols if x.id.date > self.time + timedelta(7))
            # Get contracts expires on the selected expiry
            contract_symbols = [x for x in contract_symbols if x.id.date == expiry]

            dividend_yield_model = self.securities[symbol].dividend_yield_provider

            def get_delta(contracts):
                # Mirror option contract pairs
                call, put = contracts
                if call.id.option_right == OptionRight.PUT:
                    call, put = put, call
                delta = Delta(call, self.risk_free_interest_rate_model, dividend_yield_model, put, option_model)
                self.indicator_history(delta, [call, put, call.Underlying], 1)
                return call, delta.current.value

            # Use strike price to group option contracts and get the delta
            delta_by_call = dict([get_delta(group) for strike, group in groupby(sorted(contract_symbols, key=keyfunc), keyfunc)])

            # Get the strike closest to delta of 0.75
            call, delta = sorted(delta_by_call.items(), key=lambda x: abs(x[1]-0.75))[0]

            # Let's not buy more calls
            option = self.securities.get(call)
            if option and option.invested:
                continue
            
            # Subscribe and buy one contract
            self.add_option_contract(call)
            self.market_order(call, 1, tag=f"delta={delta}")

            # Save the contracts to the underlying to liquidate and remove the subscrition when the underlying is removed 
            self.securities[symbol].contracts.append(call)