# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from AlgorithmImports import *

### <summary>
### Basic template algorithm that implements a fill model with combo orders
### <meta name="tag" content="trading and orders" />
### </summary>
class ComboOrdersFillModelAlgorithm(QCAlgorithm):
    '''Basic template algorithm that implements a fill model with combo orders'''

    def initialize(self) -> None:
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2019, 1, 20)

        self.spy = self.add_equity("SPY", Resolution.HOUR)
        self.ibm = self.add_equity("IBM", Resolution.HOUR)

        # Set the fill model
        self.spy.set_fill_model(CustomPartialFillModel())
        self.ibm.set_fill_model(CustomPartialFillModel())

        self._order_types = {}

    def on_data(self, data: Slice) -> None:
        if not self.portfolio.invested:
            legs = [Leg.create(self.spy.symbol, 1), Leg.create(self.ibm.symbol, -1)]
            self.combo_market_order(legs, 100)
            self.combo_limit_order(legs, 100, round(self.spy.bid_price))

            legs = [Leg.create(self.spy.symbol, 1, round(self.spy.bid_price) + 1), Leg.create(self.ibm.symbol, -1, round(self.ibm.bid_price) + 1)]
            self.combo_leg_limit_order(legs, 100)

    def on_order_event(self, order_event: OrderEvent) -> None:
        if order_event.status == OrderStatus.FILLED:
            order_type = self.transactions.get_order_by_id(order_event.order_id).type
            if order_type == OrderType.COMBO_MARKET and order_event.absolute_fill_quantity != 50:
                raise AssertionError(f"The absolute quantity filled for all combo market orders should be 50, but for order {order_event.order_id} was {order_event.absolute_fill_quantity}")
            elif order_type == OrderType.COMBO_LIMIT and order_event.absolute_fill_quantity != 20:
                raise AssertionError(f"The absolute quantity filled for all combo limit orders should be 20, but for order {order_event.order_id} was {order_event.absolute_fill_quantity}")
            elif order_type == OrderType.COMBO_LEG_LIMIT and order_event.absolute_fill_quantity != 10:
                raise AssertionError(f"The absolute quantity filled for all combo leg limit orders should be 10, but for order {order_event.order_id} was {order_event.absolute_fill_quantity}")

            self._order_types[order_type] = 1

    def on_end_of_algorithm(self) -> None:
        if len(self._order_types) != 3:
            raise AssertionError(f"Just 3 different types of order were submitted in this algorithm, but the amount of order types was {len(self._order_types)}")

        if OrderType.COMBO_MARKET not in self._order_types.keys():
            raise AssertionError(f"One Combo Market Order should have been submitted but it was not")

        if OrderType.COMBO_LIMIT not in self._order_types.keys():
            raise AssertionError(f"One Combo Limit Order should have been submitted but it was not")

        if OrderType.COMBO_LEG_LIMIT not in self._order_types.keys():
            raise AssertionError(f"One Combo Leg Limit Order should have been submitted but it was not")


class CustomPartialFillModel(FillModel):
    '''Implements a custom fill model that inherit from FillModel. Overrides combo_market_fill, combo_limit_fill and combo_leg_limit_fill
       methods to test FillModelPythonWrapper works as expected'''

    def __init__(self) -> None:
        self.absolute_remaining_by_order_id = {}

    def fill_orders_partially(self, parameters: FillModelParameters, fills: list[OrderEvent], quantity: int) -> list[OrderEvent]:
        partial_fills = []
        if len(fills) == 0:
            return partial_fills

        for kvp, fill in zip(sorted(parameters.securities_for_orders, key=lambda x: x.key.id), fills):
            order = kvp.key

            absolute_remaining = self.absolute_remaining_by_order_id.get(order.id, order.absolute_quantity)

            # Set the fill amount
            fill.fill_quantity = np.sign(order.quantity) * quantity

            if (min(abs(fill.fill_quantity), absolute_remaining) == absolute_remaining):
                fill.fill_quantity = np.sign(order.quantity) * absolute_remaining
                fill.status = OrderStatus.FILLED
                self.absolute_remaining_by_order_id.pop(order.id, None)
            else:
                fill.status = OrderStatus.PARTIALLY_FILLED
                self.absolute_remaining_by_order_id[order.id] = absolute_remaining - abs(fill.fill_quantity)
                price = fill.fill_price
                # self.algorithm.debug(f"{self.algorithm.time} - Partial Fill - Remaining {self.absolute_remaining_by_order_id[order.id]} Price - {price}")

            partial_fills.append(fill)

        return partial_fills

    def combo_market_fill(self, order: Order, parameters: FillModelParameters) -> list[OrderEvent]:
        fills = super().combo_market_fill(order, parameters)
        partial_fills = self.fill_orders_partially(parameters, fills, 50)
        return partial_fills

    def combo_limit_fill(self, order: Order, parameters: FillModelParameters) -> list[OrderEvent]:
        fills = super().combo_limit_fill(order, parameters)
        partial_fills = self.fill_orders_partially(parameters, fills, 20)
        return partial_fills

    def combo_leg_limit_fill(self, order: Order, parameters: FillModelParameters) -> list[OrderEvent]:
        fills = super().combo_leg_limit_fill(order, parameters)
        partial_fills = self.fill_orders_partially(parameters, fills, 10)
        return partial_fills