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

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

# region imports
from AlgorithmImports import *
from QuantConnect.Securities.Option import OptionPriceModels
from BracketOrder import BracketOrder
# endregion

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
        self.entryTicket = None
        self.stopLossTicket = None
        self.takeProfitTicket = None
    
    def placeOrder(self):
        self.entryTicket = self.algorithm.MarketOrder(self.symbol, self.totQuantity)
        self.stopLossTicket = self.algorithm.StopMarketOrder(self.symbol, -self.totQuantity, self.stopPrice)
        self.takeProfitTicket = self.algorithm.LimitOrder(self.symbol, -self.totQuantity, self.profitPrice)
    
    
    def onOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            if orderEvent.OrderId == self.entryTicket.OrderId:
              # The main order (parent sell) has been filled, place SL and TP orders
              return
            elif orderEvent.OrderId == self.stopLossTicket.OrderId or orderEvent.OrderId == self.takeProfitTicket.OrderId:
                 # Either SL or TP has been filled, cancel the opposite order
                 if self.stopLossTicket is not None and self.takeProfitTicket is not None:
                     # Make sure both orders exist before attempting to cancel
                     if orderEvent.OrderId == self.stopLossTicket.OrderId:
                          #SL order was filled, cancel TP
                          self.algorithm.Transactions.CancelOrder(self.takeProfitTicket.OrderId)
                     elif orderEvent.OrderId == self.takeProfitTicket.OrderId:
                          # TP order was filled, cancel SL
                          self.algorithm.Transactions.CancelOrder(self.stopLossTicket.OrderId)



class MyAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023, 2, 1)
        self.SetCash(1000000)

        equity = self.AddEquity("SPY", Resolution.Minute)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)

        # Add an option universe
        option = self.AddOption("SPY", Resolution.Minute)
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        option.SetFilter(-70, 0, timedelta(55), timedelta(65))

        self.symbol = option.Symbol
        self.linked_orders = {}

        # Initiate bracket order handling dictionary
        self.bracketOrders = {}
        self.tickets = []


    def OnData(self, slice):
        # Check if it's 10:00 AM
        if self.Time.hour == 10 and self.Time.minute == 0:
            option_chain = slice.OptionChains.get(self.symbol, None)

            if option_chain is None or not option_chain:
                self.Debug("No option chain data available.")
                return

            puts_in_range = [option for option in option_chain if 55 <= (option.Expiry - self.Time).days <= 65]

            if not puts_in_range:
                self.Debug("No options within the specified range of days until expiry.")
                return

            puts_sorted_by_expiry = sorted(puts_in_range, key=lambda x: abs(x.Expiry - self.Time - timedelta(days=60)))

            put_to_sell = None
            min_delta_difference = float('inf')

            for put_option in puts_sorted_by_expiry:
                delta_difference = abs(put_option.Greeks.Delta - (-0.2))

                if delta_difference < min_delta_difference:
                    put_to_sell = put_option
                    min_delta_difference = delta_difference

            if put_to_sell is not None:
                ask_price = slice.Bars[put_to_sell.Symbol].Close
                stop_loss_price = round(ask_price * 3, 2)
                take_profit_price = round(ask_price * 0.5, 2)

                # Store the BracketOrder instance in a variable
                bracket_order_instance = BracketOrder(self, put_to_sell.Symbol, -1, ask_price, ask_price, stop_loss_price, take_profit_price)

                # Use the BracketOrder class to manage the orders
                bracket_order_instance.placeOrder()

                # Store the BracketOrder instance in a list or dictionary for future reference if needed
                self.bracketOrders[put_to_sell.Symbol] = bracket_order_instance

                # Print information about the selected put option for debugging
                self.Debug(f"Selected Contract Symbol: {put_to_sell.Symbol}, " +
                        f"Strike Price: {put_to_sell.Strike}, " +
                        f"Delta: {put_to_sell.Greeks.Delta}, " +
                        f"Days Until Expiry: {(put_to_sell.Expiry - self.Time).days}")
            else:
                self.Debug("No options found that meet the criteria")