# region imports
from AlgorithmImports import *
# endregion

class CustomExerciseModelAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2022, 7, 1)
        self.set_end_date(2022, 7, 10)
        self.set_cash(100000)
        self.set_security_initializer(MySecurityInitializer(self))   
        option = self.add_option("SPY")
        self.option_symbol = option.symbol

    def on_data(self, data: Slice):
        if self.portfolio.invested:
            return
        chain = data.option_chains.get(self.option_symbol)
        if chain:
            contract = sorted([x for x in chain if x.right == OptionRight.CALL],
                key=lambda x: x.strike)[0]
            self.market_order(contract.symbol, 1)
            self.exercise_option(contract.symbol, 1)
            self.quit()

class MySecurityInitializer(BrokerageModelSecurityInitializer):

    def __init__(self, algorithm) -> None:
        super().__init__(algorithm.brokerage_model, FuncSecuritySeeder(algorithm.get_last_known_prices))
        self.algorithm = algorithm

    def initialize(self, security: Security) -> None:
        super().initialize(security)
        if security.type == SecurityType.OPTION:
            security.set_option_exercise_model(MyOptionExerciseModel(self.algorithm))

# This custom model implements the default model in LEAN (written in C#)
class MyOptionExerciseModel(DefaultExerciseModel):
    def __init__(self, algorithm):
        self.algorithm = algorithm

    def option_exercise(self, option: Option, order: OptionExerciseOrder) -> List[OrderEvent]:
        order_events = []

        underlying = option.underlying
        utc_time = Extensions.convert_to_utc(option.local_time, option.exchange.time_zone)

        in_the_money = option.is_auto_exercised(underlying.close)
        is_assignment = in_the_money and option.holdings.is_short
        messages = Messages.DefaultExerciseModel

        order_event = OrderEvent(
            order.id,
            option.symbol,
            utc_time,
            OrderStatus.FILLED,
            Extensions.get_order_direction(order.quantity),
            0,
            order.quantity,
            OrderFee.ZERO,
            messages.contract_holdings_adjustment_fill_tag(in_the_money, is_assignment, option)
        )
        order_event.is_assignment = is_assignment
        order_event.is_in_the_money = in_the_money
        order_events.append(order_event)

        if in_the_money and option.exercise_settlement == SettlementType.PHYSICAL_DELIVERY:
            exercise_quantity = option.get_exercise_quantity(order.quantity);
            order_event =  OrderEvent(
                order.id,
                underlying.symbol,
                utc_time,
                OrderStatus.FILLED,
                Extensions.get_order_direction(exercise_quantity),
                option.strike_price,
                exercise_quantity,
                OrderFee.ZERO,
                messages.option_assignment if is_assignment else messages.option_exercise
            )
            order_event.is_in_the_money = True
            order_events.append(order_event)

        return order_events