# region imports
from AlgorithmImports import *
# endregion

class CryingRedOrangeKitten(QCAlgorithm):
    def initialize(self):
        self.set_start_date(2022, 11, 22)
        self.set_end_date(2022, 12, 20)
        self.underlying = self.add_equity("SPY", extended_market_hours=True)
        self.option = self.add_option(self.underlying.symbol)
        self.option.set_filter(lambda u: u.calls_only().strikes(0,1).expiration(0,31))
        self.fixing_price = None
        
        # added for fixing extended hours issue
        self.set_security_initializer(MySecurityInitializer(self))
        self.schedule.on(self.date_rules.every_day("SPY"), self.time_rules.at(16, 0,), self.calling_fixing_price)

    # added for fixing extended hours issue
    def calling_fixing_price(self) -> None:
        
        self.fixing_price = self.underlying.close
        self.Debug(f"fixing price at {self.Time} for {self.fixing_price}")

    def on_data(self, data: Slice):
        if not self.portfolio.invested:
            chain = data.option_chains.get(self.option.symbol)
            if chain:
                self.market_order(list(chain)[0].symbol, 1)

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))
            option_exercise_model = MyOptionExerciseModel(self.algorithm)
            option_exercise_model.update_fixing_price(self.algorithm.fixing_price)
            security.set_option_exercise_model(option_exercise_model)

# This custom model implements the default model in LEAN (written in C#)
class MyOptionExerciseModel(DefaultExerciseModel):
    def __init__(self, algorithm):
        self.algorithm = algorithm
        self.fixing_price = None # added for fixing extended hours issue

    def update_fixing_price(self, price):
        self.fixing_price_during_exercise = price # added for fixing extended hours issue
        self.algorithm.Log(f"Updated fixing price to: {price}")

    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)

        # added for fixing extended hours issue
        in_the_money = option.is_auto_exercised(underlying.close)
        #in_the_money = (self.fixing_price_during_exercise >= option.strike_price) if option.right == OptionRight.CALL else (self.fixing_price_during_exercise <= option.strike_price)

        self.algorithm.Log(f"Exercising option: {option.symbol}, In-the-money: {in_the_money}, Fixing Price: {self.fixing_price}, Strike Price: {option.strike_price}, Origininal price {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