from AlgorithmImports import *
from OrderStatuses import OrderStatusCodes

'''
    Class to handle bracket orders
'''
class BracketOrder:
    def __init__(self, algorithm, parentTicket, symbol, totQuantity, stopPrice, profitPrice):
        self.algorithm = algorithm

        self.parentTicket = parentTicket
        self.symbol = symbol
        self.totQuantity = totQuantity
        self.stopPrice = stopPrice
        self.profitPrice = profitPrice

        # Initialize future properties
        self.filledQuantity = 0
        self.stopOrderTicket = None
        self.profitOrderTicket = None
    
    '''
        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"{self.algorithm.Time}, Parent ID: {self.parentId}, Order event ID: {orderEvent.OrderId} New fill 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)

            self.algorithm.Debug(f"Parent order ({self.symbol}) {OrderStatusCodes[orderEvent.Status]}, quantity: {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}")

            # Whole order filled
            if (abs(eventFillQuantity) == abs(self.filledQuantity)):
                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)
            # Connect the stop order ticket to this object
            self.algorithm.bracketOrders[self.stopOrderTicket.OrderId] = self

        # Update existing stop order ticket
        updateSettings = UpdateOrderFields()
        updateSettings.Quantity = stopQuantity
        #self.algorithm.Debug(f"Before change quantity: {self.stopOrderTicket.Quantity}")
        response = self.stopOrderTicket.Update(updateSettings)
        #self.algorithm.Debug(f"After change quantity: {self.stopOrderTicket.Quantity}")


    '''
        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)
            # Connect the stop order ticket to this object
            self.algorithm.bracketOrders[self.profitOrderTicket.OrderId] = self

        # Update existing profit order ticket
        updateSettings = UpdateOrderFields()
        updateSettings.Quantity = profitQuantity
        response = self.profitOrderTicket.Update(updateSettings)    

    '''
        Cancel trade
    '''
    def CancelTrade(self):
        # Cancel open orders and liquidate position (pre-caution)
        self.algorithm.Transactions.CancelOpenOrders(self.symbol)
        self.algorithm.Liquidate(self.symbol)

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



'''
    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

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, 3)  # 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", 10),
                        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 = {}

    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"):
                ticker = bar.Symbol

                # Create stop limit order on break of max open/close
                placementPrice = max(bar.Close, bar.Open)
                limitPrice = placementPrice + 0.05
                qty = 10
                parentTicket = self.StopLimitOrder(ticker, qty, placementPrice, limitPrice)
                self.bracketOrders[parentTicket.OrderId] = BracketOrder(self, parentTicket, parentTicket.Symbol, qty, placementPrice - 1, placementPrice + 1)
                self.Debug(f"Parent order placed for {parentTicket.Symbol}")
            # self.Debug(f"Symbol: {bar.Symbol}, Time Bar {bar.Time}, CLOSE {bar.Close}")
        

    def OnOrderEvent(self, orderEvent: OrderEvent):
        
        # 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


        self.Debug(f"Type: {orderEvent.Status}, quantity: {orderEvent.FillQuantity}")
        self.Debug(f"Bracket order keys: {self.bracketOrders.keys()}")
        self.Debug(f"Order event ID: {orderEvent.OrderId}")
        # 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.Liquidate()
        #self.Debug(f"Liquidation initiated, {self.Time}")