# region imports
from AlgorithmImports import *
# endregion

class CustomExerciseModelAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 7, 1)
        self.SetEndDate(2022, 7, 10)
        self.SetCash(100000)

        self.SetSecurityInitializer(MySecurityInitializer(self.BrokerageModel, FuncSecuritySeeder(self.GetLastKnownPrices)))

        option = self.AddOption("SPY")
        self.option_symbol = option.Symbol

    def OnData(self, data: Slice):
        chain = data.OptionChains.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.MarketOrder(contract.Symbol, 1)
            self.ExerciseOption(contract.Symbol, 1)
            self.Quit()


class MySecurityInitializer(BrokerageModelSecurityInitializer):

    def __init__(self, brokerage_model: IBrokerageModel, security_seeder: ISecuritySeeder) -> None:
        super().__init__(brokerage_model, security_seeder)

    def Initialize(self, security: Security) -> None:
        super().Initialize(security)
        if security.Type == SecurityType.Option:
            security.SetOptionExerciseModel(MyOptionExerciseModel())

class MyOptionExerciseModel(DefaultExerciseModel):
    def OptionExercise(self, option: Option, order: OptionExerciseOrder) -> List[OrderEvent]:
        in_the_money = option.IsAutoExercised(option.Underlying.Close)
        is_assignment = in_the_money and option.Holdings.IsShort

        order_event = OrderEvent(
            order.Id,
            option.Symbol,
            Extensions.ConvertToUtc(option.LocalTime, option.Exchange.TimeZone),
            OrderStatus.Filled,
            Extensions.GetOrderDirection(order.Quantity),
            0.0,
            order.Quantity,
            OrderFee.Zero,
            "MyTag"
        )
        order_event.IsAssignment = is_assignment
        return [ order_event ]