from AlgorithmImports import *
class CustomFutureFillModelAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.set_start_date(2024, 6, 3)
        self.set_end_date(2024, 6, 4)
        self.set_security_initializer(CustomSecurityInitializer(self))
        self.es = self.add_future("ES", extended_market_hours=True)
        self.ticket = None
    def on_data(self, slice):
        if not self.ticket:
            self.market_order(self.es.mapped, 1)
            self.ticket = self.trailing_stop_order(self.es.mapped, -1, 5, False)
        if self.ticket and self.ticket.status != OrderStatus.FILLED:
            self.plot(self.es.mapped.value, 'Price', self.es.price)
            self.plot(self.es.mapped.value, 'Stop', self.ticket.get(OrderField.STOP_PRICE))

class CustomSecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, algorithm):
        super().__init__(algorithm.brokerage_model, SecuritySeeder.NULL)
    def initialize(self, security):
        super().initialize(security)
        security.set_fill_model(CustomFutureFillModel())

class CustomFutureFillModel(FutureFillModel):
    def trailing_stop_fill(self, asset, order):
        utc_time = Extensions.convert_to_utc(asset.local_time, asset.exchange.time_zone)
        fill = OrderEvent(order, utc_time, OrderFee.ZERO)
        if order.Status == OrderStatus.CANCELED:
            return fill

        bar = asset.cache.get_data[TradeBar]()

        # Do not fill on stale data
        if Extensions.convert_to_utc(bar.end_time, asset.exchange.time_zone) <= order.time:
            return fill

        if order.direction == OrderDirection.BUY and bar.high >= order.stop_price:
            fill.Status = OrderStatus.FILLED
            fill.FillQuantity = order.quantity
            # Assuming worse case scenario fill - fill at highest of the stop & asset bid price.
            slip = asset.SlippageModel.GetSlippageApproximation(asset, order)
            fill.FillPrice = max(order.stop_price, asset.ask_price + slip)
        if order.direction == OrderDirection.SELL and bar.low <= order.stop_price:
            fill.Status = OrderStatus.FILLED
            fill.FillQuantity = order.quantity
            # Assuming worse case scenario fill - fill at lowest of the stop & asset bid price.
            slip = asset.SlippageModel.GetSlippageApproximation(asset, order)
            fill.FillPrice = min(order.stop_price, asset.bid_price - slip)

        if fill.Status != OrderStatus.FILLED:
            current_market_price = bar.high if order.direction == OrderDirection.SELL else bar.low
            is_update, updated_stop_price = TrailingStopOrder.try_update_stop_price(current_market_price, order.stop_price, order.trailing_amount, order.trailing_as_percentage, order.Direction, 0)
            if is_update:
                order.stop_price = updated_stop_price
                self.parameters.on_order_updated(order)

        return fill