from AlgorithmImports import *
class CustomDataEMFAlgorithm(QCAlgorithm):

    def initialize(self):

        self.set_start_date(2019, 1, 1)
        self.set_end_date(2024, 2, 1)
        self.set_cash(10000000)
        self.activeSecurities = set()

        self.Settings.FreePortfolioValuePercentage = 0.05

        ### Set benchmark
        # self.benchmark_symbol = self.get_parameter("benchmark_ticker")
        self.set_benchmark("EEM")
        self.AddEquity("EEM", Resolution.HOUR)

        universe = self.add_universe(MyCustomUniverseDataClass, "myCustomUniverse", self.selector_function)
        self.universe_settings.resolution = Resolution.DAILY

        self.portfolioTargets = []
        self.aro_data = {1: {}, 3: {}, 5: {}} # Add this line to initialize the dictionary
        self.cef_discount_data = {} 
        self.cef_purchase_price = {}
        self.cef_purchase_discount = {}
        self.cef_purchase_percentile = {}

        ### Helper variables
        self.take_profit_active = {} #Is take profit flag activated?
        self.profit_threshold = {}

        ### Other Helper variables
        self.firstOrderTicket = {}
        self.firstOrderTicketFillFlag = {}
        self.takeProfitMarketTicket = {}
        self.firstOrderEntryTime = {}
        self.firstOrderEntryTicketFillTime = datetime.min
        self.dipsTimeCheck = {}
        self.invalidOrderReSubmit = {}
        self.invalidOrderCount = 0

        self.averageFillPrices = {}
        self.buyDipsMarketTicket = {}

        ### Define take profit parameters - also used in trimming
        self.take_profit_percent = float(self.GetParameter("take_profit_percent"))
        self.price_rise_percent = float(self.GetParameter("price_rise_percent"))

        ### Define discount percentile upper and lower limits below
        self.aro_dptl_lower_limit = float(self.get_parameter("dptl_lower_limit"))
        self.aro_dptl_upper_limit = float(self.get_parameter("dptl_upper_limit"))

        ### Define discount differential 
        self.discount_differential = float(self.get_parameter("discount_differential_limit"))

        ###Define period whether, 1, 3, or 5 year below
        self.period = int(self.get_parameter("period_selection"))

     # Define the selector function
    def selector_function(self, data):
        sorted_data = sorted([ x for x in data if x["CEF_Discount"] <= 4 ],
                            key=lambda x: x.end_time) #,
                            # reverse=True)
        self.Log('Stocks trading at a discount to NAV: ' + str(len(sorted_data)))
        for x in sorted_data:
            # self.aro_data[5][x.symbol] = x["ARO_DPTL_5YR"]
            # self.aro_data[3][x.symbol] = x["ARO_DPTL_3YR"]
            self.aro_data[1][x.symbol] = x["ARO_DPTL_1YR"]
            self.cef_discount_data[x.symbol] = x["CEF_Discount"]
        return [x.symbol for x in sorted_data]    

    def on_securities_changed(self, changes):
        # remove stocks from active secutiries positions
        for x in changes.RemovedSecurities:
            if x.Symbol in self.activeSecurities:
                self.activeSecurities.remove(x.Symbol)
                         

        # can't open positions here since data might not be added correctly yet
        for x in changes.AddedSecurities:
            # Skip if the added security is the benchmark
            if x.Symbol.Value == "EEM":
                continue
            self.Debug(f"{self.Time}: Added {x.Symbol}")
            self.activeSecurities.add(x.Symbol)
            self.Securities[x.Symbol].SetLeverage(1.0)
            

        # adjust targets if universe has changed
        self.portfolioTargets = [PortfolioTarget(symbol, 1/len(self.activeSecurities)) 
                            for symbol in self.activeSecurities]
        
    def OnData(self, data):



        # check whether 30 days passed since 1st buy order before selling
        
        if self.portfolioTargets == []:
            return
            
        for target in self.portfolioTargets:
            symbol = target.Symbol
            price = round(self.Securities[symbol].Price,2)  # Get the current price
            aro_value = round(self.aro_data[self.period][symbol],2)
            cef_discount_value = self.cef_discount_data[symbol] # Get the CEF discount value
           
            ###Logic for topping-up
            # Step 1: Calculate buying power and order size
            buyingPower = self.Portfolio.MarginRemaining
            price = self.Securities[symbol].Price
            if price != 0:
                orderSize = max(1, buyingPower / price)
            else:
                self.Debug(f"{symbol} Price is ({price}) on {self.time}, cannot calculate order size")
                return

            # Step 2: Check top-up conditions are met
            if self.Portfolio[symbol].Invested and not self.transactions.get_open_orders(symbol) \
                and self.Portfolio[symbol].HoldingsValue < self.Portfolio.TotalPortfolioValue * 0.1 \
                    and aro_value < self.aro_dptl_lower_limit \
                    and (abs(cef_discount_value - self.cef_purchase_discount[symbol]) > self.discount_differential or price < self.Portfolio[symbol].average_price): # and (self.Time - self.dipsTimeCheck[symbol]).days >= 1:
                
                # self.Debug(f"time: {self.time}")
                # self.Debug(f"symbol: {symbol}")
                # self.Debug(f"self.Portfolio[symbol].Invested: {self.Portfolio[symbol].Invested}")
                # self.Debug(f"self.firstOrderTicketFillFlag[symbol]: {self.firstOrderTicketFillFlag[symbol]}")
                # self.Debug(f"self.Portfolio[symbol].HoldingsValue: {self.Portfolio[symbol].HoldingsValue}")
                # self.Debug(f"self.Portfolio.TotalPortfolioValue: {self.Portfolio.TotalPortfolioValue}")
                # self.Debug(f"cef_discount_value: {cef_discount_value}")
                # self.Debug(f"self.Portfolio[symbol].average_price: {self.Portfolio[symbol].average_price}")
                # self.Debug(f"self.cef_purchase_discount: {self.cef_purchase_discount[symbol]}")
                
                # Step 3: Adjust the order size according to the available buying power
                adjusted_order_size = min(orderSize, self.Portfolio[symbol].Quantity * self.take_profit_percent)

                self.buyDipsMarketTicket[symbol] = self.limit_order(symbol, adjusted_order_size, \
                    price,tag=f"Topping up {symbol} @ Discount: {cef_discount_value} / {self.cef_purchase_discount[symbol]} @ ARO_DPTL_{self.period}YR: {aro_value} @ Price: {price} / {round(self.Portfolio[symbol].average_price,2)}" )
                # self.dipsTimeCheck[symbol] = self.Time

            ###Logic for trimming positions
            if self.Portfolio[symbol].Invested and not self.transactions.get_open_orders(symbol) \
                and aro_value > self.aro_dptl_upper_limit \
                    and abs(cef_discount_value - self.cef_purchase_discount[symbol]) > self.discount_differential: # or (price / self.averageFillPrices[symbol] > 1+self.price_rise_percent):

                orderSize = min(-1, -self.Portfolio[symbol].Quantity * self.take_profit_percent)
                self.takeProfitMarketTicket[symbol] = self.limit_order(symbol, orderSize, \
                    price,tag=f"Trimming {symbol} @ Discount: {cef_discount_value} / {self.cef_purchase_discount[symbol]} @ ARO_DPTL_{self.period}YR: {aro_value} @ Price: {price} / {self.cef_purchase_price[symbol]}")
    
            ### INITIATE BUY
            if not self.Portfolio[symbol].Invested and not self.transactions.get_open_orders(symbol) and aro_value < self.aro_dptl_lower_limit and cef_discount_value < 0:
                self.cef_purchase_discount[symbol] = cef_discount_value
                self.cef_purchase_price[symbol] = price
                self.cef_purchase_percentile[symbol] = aro_value
                self.SetHoldings([target], tag=f"Initiate BUY {symbol} @ Discount: {cef_discount_value} @ ARO_DPTL_{self.period}YR: {aro_value}") # Buy
                self.log(f"Initiate BUY {symbol} @ Discount: {cef_discount_value} @ ARO_DPTL_{self.period}YR: {aro_value}") # Buy
                self.firstOrderEntryTime[symbol] = self.time

            ### RE-INITIATE INVALID BUY INVALID DUE TO INSUFFICIENT MARGIN
            if symbol in self.invalidOrderReSubmit and not self.Portfolio[symbol].Invested and self.invalidOrderReSubmit[symbol].Status == OrderStatus.INVALID:
                self.Debug(f"Im in ondata and invalidOrderReSubmit {symbol}: {self.invalidOrderReSubmit[symbol]}")
                # Calculate the available buying power and order size and only if > 0 proceed
                buyingPower = self.Portfolio.MarginRemaining
                self.Debug(f"self.Portfolio.MarginRemaining: {self.Portfolio.MarginRemaining}")
                if buyingPower > 0 and self.invalidOrderCount < 2 :
                    self.Debug(f"Ondata Counter: {self.invalidOrderCount}")
                    price = self.Securities[symbol].Price
                    self.Debug(f"price after margin check: {price}")
                    orderSize = max(1, buyingPower / price)
                    # Calculate the total value of the order 
                    total_order_value = price * orderSize
                    ten_percent_portfolio = 1 * self.Portfolio.TotalPortfolioValue
                    # Check if the total order value is greater than 10% of the portfolio value
                    if total_order_value > ten_percent_portfolio:
                        # Adjust the order size
                        orderSize = ten_percent_portfolio / price
                        self.Debug(f"Adjusted order size to maintain 10% portfolio value limit: {orderSize}")
                    self.Debug(f"ordersize after margin check: {orderSize}")
                # Re-submit invalid order with new size and clear 
                    self.cef_purchase_discount[symbol] = cef_discount_value
                    self.cef_purchase_price[symbol] = price
                    self.cef_purchase_percentile[symbol] = aro_value
                    self.market_order(symbol, orderSize, tag=f"Initiate BUY re-submit w/lower size @ Discount: {cef_discount_value} @ ARO_DPTL_{self.period}YR: {aro_value}")
                    self.Debug(f"Initiate BUY re-submit w/lower size ({orderSize}) @ Discount: {cef_discount_value} @ ARO_DPTL_{self.period}YR: {aro_value}")
                    self.firstOrderEntryTime[symbol] = self.time
                    # del self.invalidOrderReSubmit[symbol]
                    

            elif (self.time - self.firstOrderEntryTicketFillTime).days >= 30 and cef_discount_value > 0 and self.Portfolio[symbol].Invested and abs(cef_discount_value - self.cef_purchase_discount[symbol]) > self.discount_differential and (price / self.cef_purchase_price[symbol]) > 1:
                if self.Portfolio[symbol].Quantity <= 0:
                    return
                self.Debug(f"Liquidating {symbol} at {self.Time}, First order entry time: {self.firstOrderEntryTicketFillTime}, Days passed = {(self.time - self.firstOrderEntryTicketFillTime).days}, Price: {price}, purchase price @ {self.cef_purchase_price[symbol]} ARO_DPTL_{self.period}YR: {aro_value}, Selling Discount: {cef_discount_value}, Purchased at Discount: {self.cef_purchase_discount[symbol]}")  
                self.Liquidate(symbol, f"Liquidating {symbol} @ Discount: {cef_discount_value} - Differential: {cef_discount_value - self.cef_purchase_discount[symbol]}") # Sell
                self.take_profit_active[symbol] = False

            self.Plot("ARO Data", f"ARO_DPTL_{self.period}YR", self.aro_data[self.period][symbol])
            self.Plot("ARO Data", "CEF_Discount", cef_discount_value)
            self.Plot("ARO Data", "Price", price)
            benchmark_value = self.Benchmark.Evaluate(self.Time)
            self.Plot("ARO Data", "Benchmark:EWL", benchmark_value)
        
        
        # self.portfolioTargets = []

    def on_order_event(self, order_event):

        # order = self.Transactions.GetOrderById(order_event.OrderId)
        # self.Log(f"Order details: Id={order.OrderId}, Symbol={order.Symbol}, Type={order.Type}, Quantity={order.Quantity}, Time={order.Time}")
        
        entryTicket = self.Transactions.GetOrderTicket(order_event.OrderId)

        symbol = order_event.Symbol
        cef_discount_value = self.cef_discount_data[symbol]
        aro_value = round(self.aro_data[self.period][symbol],2)
        current_weighted_discount = self.Portfolio[symbol].Quantity * self.cef_purchase_discount[symbol]
        self.averageFillPrices[symbol] = round(entryTicket.AverageFillPrice,2)

        self.firstOrderTicket[symbol] = entryTicket
        self.buyDipsMarketTicket[symbol] = entryTicket

        if order_event.Status == OrderStatus.INVALID:
            self.invalidOrderCount += 1
            if self.invalidOrderCount == 2:
                return    
                # self.Debug(str(order_event))
            self.invalidOrderReSubmit[symbol] = entryTicket
                
   
        if order_event.Status == OrderStatus.FILLED:
            ### Change the flag to true if its the first order filled!
            if "Initiate BUY" in entryTicket.Tag:
                self.firstOrderTicketFillFlag[symbol] = True
            ### Change the flag to false if a symbol order is filled and no longer held!
            if "Liquidating" in entryTicket.Tag:
                self.firstOrderTicketFillFlag[symbol] = False
            ### Logic for updated CEF Purchase Discount once a new top-up order has been filled!
            if "Topping up" in entryTicket.Tag:
                new_weighted_discount = ((self.Portfolio[symbol].Quantity * self.take_profit_percent / 2) * cef_discount_value)
                new_quantity = self.Portfolio[symbol].Quantity + (self.Portfolio[symbol].Quantity * self.take_profit_percent / 2)
                new_average_purchase_discount = (current_weighted_discount + new_weighted_discount) / new_quantity
                self.cef_purchase_discount[symbol] = round(new_average_purchase_discount, 2)

        # save fill time of initial buy order and wait 30 days before 1st sale
        if self.firstOrderTicket[symbol] is not None and self.firstOrderTicket[symbol].order_id == order_event.order_id:
            self.firstOrderEntryTicketFillTime = self.time

        if order_event.status != OrderStatus.FILLED:
            return

    def OnEndOfDay(self):
        self.invalidOrderCount = 0


# Example custom universe data; it is virtually identical to other custom data types.
class MyCustomUniverseDataClass(PythonData):		
    def get_source(self, config, date, is_live_mode):
        source = "https://raw.githubusercontent.com/jabertech/backtest-qc/main/combined_selection_3.csv"
        return SubscriptionDataSource(source, SubscriptionTransportMedium.RemoteFile);
		
    def Reader(self, config, line, date, isLive):
        if not (line.strip() and line[0].isalnum()):
            return None
			
        items = line.split(",")
    
        # Generate required data, then return an instance of your class.
        data = MyCustomUniverseDataClass()
		
        try:
            data.end_time = datetime.strptime(items[1], '%Y-%m-%d')+timedelta(hours=20)

            # define Time as exactly 1 day earlier Time
            data.time = data.end_time - timedelta(1)

            data.symbol = Symbol.create(items[0], SecurityType.EQUITY, Market.USA)
            data.value = float(items[5])
            data["Open"] = float(items[2])
            data["High"] = float(items[3])
            data["Low"] = float(items[4])
            data["Close"] = float(items[5])
            data["Volume"] = int(items[6])
            data["CEF_NAV"] = float(items[7])
            data["CEF_Discount"] = float(items[8])
            data["CEF_Price"] = float(items[9])
            data["ARO_DPTL_1YR"] = float(items[10])
            # data["ARO_DPTL_3YR"] = float(items[12])
            # data["ARO_DPTL_5YR"] = float(items[14])
        except ValueError:
            return None
        
        return data