# region imports
from AlgorithmImports import *
# endregion

class CalculatingFluorescentOrangeSeahorse(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2024, 4, 11)
        self.set_end_date(2024, 4, 12)

        underlying = self.add_equity("AFL").symbol

        get_symbol = lambda s: Symbol.create_option(underlying, underlying.id.market,
            OptionStyle.AMERICAN, OptionRight.CALL, s, datetime(2024,4,26))
            
        self.legs = [ Leg.create(get_symbol(80), 1), Leg.create(get_symbol(81), -1) ]

        [self.add_option_contract(leg.symbol).set_fill_model(MidPriceFillModel(self))
            for leg in self.legs]

        self.schedule.on(self.date_rules.on(self.start_date), self.time_rules.at(15,45), self._trade)

    def _trade(self):
        # Place one order with a limit of 0.5
        # The expected spread is 0.45, so it should fill
        self.combo_limit_order(self.legs, 1, .5)

    def on_order_event(self, order_event):
        self.log(str(order_event))

class MidPriceFillModel(FillModel):
    def __init__(self, algorithm):
        super().__init__()
        self._max_spread = 1
        self.algorithm = algorithm
    
    def combo_limit_fill(self, order, parameters):
        prices = []
        for kvp in parameters.SecuritiesForOrders:
            order = kvp.key
            asset = kvp.Value
            quantity = GroupOrderExtensions.get_order_leg_ratio(order.quantity, order.group_order_manager)
            prices.append((asset.bid_price + asset.ask_price) / 2 * quantity)
        
        # The discount logic will create fill prices that observe the maximum spread
        # for example
        # AFL 240426C00080000 at 1.380
        # AFL 240426C00081000 at 0.925
        # has a spread of 0.45
        # it will become 
        # AFL 240426C00080000 at 0.920
        # AFL 240426C00081000 at 0.620
        # for a self._max_spread of 0.3
        discount = 1
        current_price = sum(prices)
        if current_price > self._max_spread:
            discount = self._max_spread / current_price 
            current_price = self._max_spread

        self.algorithm.log(f'The spread is {current_price=}')
        limit_price = order.group_order_manager.limit_price
        fills = []

        def create_fill(order, asset):
            utc_time = Extensions.convert_to_utc(asset.local_time, asset.exchange.time_zone)
            fill = OrderEvent(order, utc_time, OrderFee.ZERO)
            fill.status = OrderStatus.FILLED
            mid_price = (asset.bid_price + asset.ask_price) / 2
            fill.fill_price = mid_price * discount 
            fill.fill_quantity = order.quantity
            return fill

        match order.group_order_manager.direction:
            case OrderDirection.BUY:
                if current_price < limit_price:
                    for kvp in parameters.SecuritiesForOrders:
                        fills.append(create_fill(kvp.key, kvp.Value))
            case OrderDirection.SELL:
                if current_price > limit_price:
                    for kvp in parameters.SecuritiesForOrders:
                        fills.append(create_fill(kvp.key, kvp.Value))
        return fills