from AlgorithmImports import *
from OrderStatuses import OrderStatusCodes

'''
    Class to handle bracket orders

    TODO: CREATE STOP LOSS AND PROFIT ORDERS BASED ON THE ACTUAL FILL PRICE 
'''
class BracketOrder:
    def __init__(self, algorithm, symbol, totQuantity, placementPrice, limitPrice, stopPrice, profitPrice):
        
        self.algorithm = algorithm
        self.symbol = symbol
        self.totQuantity = totQuantity
        self.placementPrice = placementPrice
        self.limitPrice = limitPrice
        self.stopPrice = stopPrice
        self.profitPrice = profitPrice

        # Initialize future properties
        self.filledQuantity = 0
        self.parentTicket = None
        self.stopOrderTicket = None
        self.profitOrderTicket = None

        # Place initial order
        self.PlaceInitialOrder()

    '''
        Place initial order
    '''
    def PlaceInitialOrder(self):
        self.parentTicket = self.algorithm.StopLimitOrder(self.symbol, self.totQuantity, self.placementPrice, self.limitPrice, tag="Initial order")
        self.algorithm.bracketOrders[self.parentTicket.OrderId] = self
        self.algorithm.Debug(f"Parent order placed for {self.symbol}")
        return      

    '''
        Handle order changes

        Orderstatus enum:
        https://www.quantconnect.com/docs/algorithm-reference/trading-and-orders#Trading-and-Orders-Tracking-Order-Events
    '''
    def HandleOrderChange(self, orderEvent):

        # If the orderstatus is PartiallyFilled or Filled it will have a filledQuantity > 0
        # Currently make no changes for order statuses that does not (partially) fill an order
        eventFillQuantity = orderEvent.FillQuantity

        if (orderEvent.FillQuantity == 0):
            return
        
        '''
            Handle scenario where parent order is updated

            Add/update stop and profit orders
        '''
        if (orderEvent.OrderId == self.parentTicket.OrderId):
            self.algorithm.Debug(f"Parent order ({self.symbol}) {OrderStatusCodes[orderEvent.Status]}, quantity: {eventFillQuantity}")
            
            self.filledQuantity += eventFillQuantity

            # Place/add to stop order (negative because we are selling)
            self.PlaceStopOrder(-eventFillQuantity)

            # Place/add to profit order (negative because we are selling)
            self.PlaceProfitOrder(-eventFillQuantity)
            return
        
        '''
            Handle scenario where stop loss is hit
        
            - Cancel all open orders on the symbol
            - Liquidate portfolio of that stock (should not be necessary, but extra precaution)
        '''
        if (orderEvent.OrderId == self.stopOrderTicket.OrderId):
            self.CancelTrade()
            self.algorithm.Debug(f"Stop order ({self.symbol}) {OrderStatusCodes[orderEvent.Status]}, quantity: {eventFillQuantity}")
            return
            

        '''
            Handle scenario where profit target is hit

            - Cancel all open orders on the symbol (for now)
            - Should in reality check if profit is partially filled (not equal to self.filledQuantity), and update stop loss order
            - Should probably cancel parent order in case it is partially filled
        '''
        if (orderEvent.OrderId == self.profitOrderTicket.OrderId):
            self.algorithm.Debug(f"Profit order ({self.symbol}) {OrderStatusCodes[orderEvent.Status]}, quantity: {eventFillQuantity}")
            self.CancelTrade()
            return
    
    '''
        Place/add to stop market order when the parent order is (partially) filled
    '''
    def PlaceStopOrder(self, stopQuantity):
        # Create new profit order ticket if non exists
        if (self.stopOrderTicket is None):
            self.stopOrderTicket = self.algorithm.StopMarketOrder(self.symbol, stopQuantity, self.stopPrice, tag=f"Stop order ({self.stopPrice})")
            # Connect the stop order ticket to this object
            self.algorithm.bracketOrders[self.stopOrderTicket.OrderId] = self
            self.algorithm.Debug(f"Stop order placed ({self.symbol})")
            return

        # Update existing stop order ticket
        updateSettings = UpdateOrderFields()
        updateSettings.Quantity = stopQuantity
        response = self.stopOrderTicket.Update(updateSettings)
        self.algorithm.Debug(f"Stop order updated ({self.symbol})")
        return

    '''
        Place/add to profit taking order when the parent order is (partially) filled
    '''
    def PlaceProfitOrder(self, profitQuantity):
        # Create new profit order ticket if non exists
        if (self.profitOrderTicket is None):
            self.profitOrderTicket = self.algorithm.LimitOrder(self.symbol, profitQuantity, self.profitPrice, tag=f"Profit order ({self.profitPrice})")
            # Connect the stop order ticket to this object
            self.algorithm.bracketOrders[self.profitOrderTicket.OrderId] = self
            self.algorithm.Debug(f"Profit order placed ({self.symbol})")
            return

        # Update existing profit order ticket
        updateSettings = UpdateOrderFields()
        updateSettings.Quantity = profitQuantity
        response = self.profitOrderTicket.Update(updateSettings)
        self.algorithm.Debug(f"Profit order updated ({self.symbol})")    
        return

    '''
        Cancel trade
    '''
    def CancelTrade(self):
        # Cancel open orders and liquidate position (pre-caution)
        self.algorithm.Transactions.CancelOpenOrders(self.symbol)
        self.algorithm.Liquidate(self.symbol, tag="LIQUIDATION ORDER INSIDE")

        # Remove order IDs from bracketOrders dictionary
        # self.algorithm.bracketOrders.pop(self.parentTicket.OrderId, None)
        # self.algorithm.bracketOrders.pop(self.profitOrderTicket.OrderId, None)
        # self.algorithm.bracketOrders.pop(self.stopOrderTicket.OrderId, None)
        return

#region imports
from AlgorithmImports import *
#endregion


OrderStatusCodes = {
    0:'NEW', # new order pre-submission to the order processor
    1:'SUBMITTED', # order submitted to the market
    2:'PARTIALLY FILLED', # partially filled, in market order
    3:'FILLED', # completed, filled, in market order
    5:'CANCELED', # order cancelled before filled
    6:'NONE', # no order state yet
    7:'INVALID', # order invalidated before it hit the market (e.g. insufficient capital)
    8:'CANCEL PENDING', # order waiting for confirmation of cancellation
    9:'UPDATE SUBMITTED' # Order update submitted to the market
}

from AlgorithmImports import *
from BracketOrder import BracketOrder

class TestBracketOrder(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2022, 6, 1)  # Set Start Date
        self.SetEndDate(2022, 6, 1)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        # Ensure securities are set with adjusted data normalization
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Adjusted))

        # Add SPY for scheduled events
        self.AddEquity("SPY", resolution=Resolution.Minute, extendedMarketHours=True)

        # Create scheduled event to liquidate open orders and positions at end of day
        self.Schedule.On(self.DateRules.EveryDay("SPY"), 
                        self.TimeRules.BeforeMarketClose("SPY", 30),
                        self.EndOfDay)

        # Add testing assets
        self.assets = ["GOOG", "AMZN"]
        for asset in self.assets:
            self.AddEquity(asset, Resolution.Minute, extendedMarketHours=True)
        
        # Initiate bracket order handling dictionary
        self.bracketOrders = {}
        self.tickets = []
        self.liquidatedOrderIds = []
        self.currentLiquidated = None

    def OnData(self, data: Slice):
        """OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        """

        for bar in data.Bars.Values:
            if (bar.Time.hour == 9 and bar.Time.minute == 30 and bar.Symbol != "SPY"):
                symbol = bar.Symbol

                # Create stop limit order on break of max open/close
                placementPrice = max(bar.Close, bar.Open)
                limitPrice = placementPrice + 0.05
                qty = 100
                # Place order - initiate bracket order object
                BracketOrder(self, symbol, qty, placementPrice, limitPrice, placementPrice * 0.98, placementPrice * 1.02)
            # self.Debug(f"Symbol: {bar.Symbol}, Time Bar {bar.Time}, CLOSE {bar.Close}")
        

    def OnOrderEvent(self, orderEvent: OrderEvent):

        # If portfolio liquidation happens from scheduled event at the end of the day, no bracketorder object is created and can not be found in the bracketorders dictionary 
        # Just ignore liquidation orders - they are tagged by "LIQUIDATION ORDER"
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if ("LIQUIDATION ORDER" in order.Tag):
            return 
        
        # If the order status is NEW or SUBMITTED, no changes are made to the order
        if (orderEvent.Status == OrderStatus.New or orderEvent.Status == OrderStatus.Submitted):
            return

        # For any other order statuses we expect the order to be in the dictionary of orders    
        bracketOrder = self.bracketOrders[orderEvent.OrderId]
        # Make changes in the bracket order object 
        bracketOrder.HandleOrderChange(orderEvent)


    '''
        Liquidate portfolio and open orders at end of day
    '''
    def EndOfDay(self):
        self.Transactions.CancelOpenOrders()
        self.currentLiquidated = self.Liquidate(tag="LIQUIDATION ORDER")
        self.Debug(f"Liquidation performed: {self.currentLiquidated}")
        

#region imports
from AlgorithmImports import *
#endregion
'''
    Random stuff

if (orderEvent.OrderId not in self.bracketOrders):
    if (orderEvent.FillQuantity != 0):
        self.Debug(f"{self.Time}, Order quantity not zero: {orderEvent.FillQuantity}")
        assert (1 == 2), "Order quantity not zero"
    else:
        return


// Place your key bindings in this file to override the defaultsauto[]
[
    {
        "key": "ctrl+k",
        "command": "editor.action.commentLine",
        "when": "editorTextFocus && !editorReadonly"
    },
    {
        "key": "ctrl+/",
        "command": "-editor.action.commentLine",
        "when": "editorTextFocus && !editorReadonly"
    }
]

# Order partially filled - cancel remaining parent order, update stop order
self.filledQuantity += eventFillQuantity
cancelRes = self.parentTicket.Cancel("Canceled remaining parent order")
# Update existing stop order ticket
updateSettings = UpdateOrderFields()
updateSettings.Quantity = self.filledQuantity
updateRes = self.stopOrderTicket.Update(updateSettings)
return

'''