### 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 *

class RegressionTestShortableProvider(LocalDiskShortableProvider):
    def __init__(self):
        super().__init__("testbrokerage")

### <summary>
### Tests that orders are denied if they exceed the max shortable quantity.
### </summary>
class ShortableProviderOrdersRejectedRegressionAlgorithm(QCAlgorithm):
    def initialize(self):
        self.orders_allowed = []
        self.orders_denied = []
        self.initialized = False
        self.invalidated_allowed_order = False
        self.invalidated_new_order_with_portfolio_holdings = False

        self.set_start_date(2013, 10, 4)
        self.set_end_date(2013, 10, 11)
        self.set_cash(10000000)

        self.spy = self.add_equity("SPY", Resolution.MINUTE)
        self.aig = self.add_equity("AIG", Resolution.MINUTE)

        self.spy.set_shortable_provider(RegressionTestShortableProvider())
        self.aig.set_shortable_provider(RegressionTestShortableProvider())

    def on_data(self, data):
        if not self.initialized:
            self.handle_order(self.limit_order(self.spy.symbol, -1001, 10000)) # Should be canceled, exceeds the max shortable quantity
            order_ticket = self.limit_order(self.spy.symbol, -1000, 10000)
            self.handle_order(order_ticket) # Allowed, orders at or below 1000 should be accepted
            self.handle_order(self.limit_order(self.spy.symbol, -10, 0.01)) # Should be canceled, the total quantity we would be short would exceed the max shortable quantity.

            response = order_ticket.update_quantity(-999) # should be allowed, we are reducing the quantity we want to short
            if not response.is_success:
                raise ValueError("Order update should of succeeded!")

            self.initialized = True
            return

        if not self.invalidated_allowed_order:
            if len(self.orders_allowed) != 1:
                raise AssertionError(f"Expected 1 successful order, found: {len(self.orders_allowed)}")
            if len(self.orders_denied) != 2:
                raise AssertionError(f"Expected 2 failed orders, found: {len(self.orders_denied)}")

            allowed_order = self.orders_allowed[0]
            order_update = UpdateOrderFields()
            order_update.limit_price = 0.01
            order_update.quantity = -1001
            order_update.tag = "Testing updating and exceeding maximum quantity"

            response = allowed_order.update(order_update)
            if response.error_code != OrderResponseErrorCode.EXCEEDS_SHORTABLE_QUANTITY:
                raise AssertionError(f"Expected order to fail due to exceeded shortable quantity, found: {response.error_code}")

            cancel_response = allowed_order.cancel()
            if cancel_response.is_error:
                raise AssertionError("Expected to be able to cancel open order after bad qty update")

            self.invalidated_allowed_order = True
            self.orders_denied.clear()
            self.orders_allowed.clear()
            return

        if not self.invalidated_new_order_with_portfolio_holdings:
            self.handle_order(self.market_order(self.spy.symbol, -1000)) # Should succeed, no holdings and no open orders to stop this
            spy_shares = self.portfolio[self.spy.symbol].quantity
            if spy_shares != -1000:
                raise AssertionError(f"Expected -1000 shares in portfolio, found: {spy_shares}")

            self.handle_order(self.limit_order(self.spy.symbol, -1, 0.01)) # Should fail, portfolio holdings are at the max shortable quantity.
            if len(self.orders_denied) != 1:
                raise AssertionError(f"Expected limit order to fail due to existing holdings, but found {len(self.orders_denied)} failures")

            self.orders_allowed.clear()
            self.orders_denied.clear()

            self.handle_order(self.market_order(self.aig.symbol, -1001))
            if len(self.orders_allowed) != 1:
                raise AssertionError(f"Expected market order of -1001 BAC to not fail")

            self.invalidated_new_order_with_portfolio_holdings = True

    def handle_order(self, order_ticket):
        if order_ticket.submit_request.status == OrderRequestStatus.ERROR:
            self.orders_denied.append(order_ticket)
            return

        self.orders_allowed.append(order_ticket)