from AlgorithmImports import *
class CustomDataEMFAlgorithm(QCAlgorithm):

    def initialize(self):

        self.set_start_date(2008, 1, 8)
        self.set_end_date(2014, 7, 25)
        self.set_cash(100000)
        self.activeSecurities = set()

        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_discount = {}

        ### 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 period whether, 1, 3, or 5 year below
        self.period = 1


     # Define the selector function
    def selector_function(self, data):
        sorted_data = sorted([ x for x in data if x["CEF_Discount"] <= 0 ],
                            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):
        # close positions in removed securities
        for x in changes.RemovedSecurities:
            self.Debug(f"{self.Time}: Removed {x.Symbol}")
            # self.Debug(f"ARO_DPTL_{self.period}YR for {x.Symbol} when removed: {self.aro_data[self.period][x.Symbol]}")
            self.Liquidate(x.Symbol)
            if x.Symbol in self.activeSecurities:
                self.Debug(f"ARO_DPTL_{self.period}YR for {x.Symbol} when removed: {self.aro_data[self.period][x.Symbol]}")
                self.activeSecurities.remove(x.Symbol)
                del self.aro_data[self.period][x.Symbol]

            

        # can't open positions here since data might not be added correctly yet
        for x in changes.AddedSecurities:
            self.Debug(f"{self.Time}: Added {x.Symbol}")
            self.Debug(f"ARO_DPTL_{self.period}YR for {x.Symbol} when added: {self.aro_data[self.period][x.Symbol]}")
            self.activeSecurities.add(x.Symbol)
            

        # adjust targets if universe has changed
        self.portfolioTargets = [PortfolioTarget(symbol, 1/len(self.activeSecurities)) 
                            for symbol in self.activeSecurities]
        
    def OnData(self, data):
        
        
        if self.portfolioTargets == []:
            return
        
        for target in self.portfolioTargets:
            symbol = target.Symbol
            price = self.Securities[symbol].Price  # Get the current price
            aro_value = self.aro_data[self.period][symbol]
            cef_discount_value = self.cef_discount_data[symbol] # Get the CEF discount value
            if aro_value < self.aro_dptl_lower_limit and not self.Portfolio[symbol].Invested:
                self.cef_purchase_discount[symbol] = self.cef_discount_data[symbol]
                self.SetHoldings(self.portfolioTargets) # Buy
                self.Debug(f"Buying {symbol} at {self.Time}, Price: {price}, ARO_DPTL_{self.period}YR: {aro_value}, Purchase Discount: {self.cef_purchase_discount[symbol]}")
            elif aro_value > self.aro_dptl_upper_limit and self.Portfolio[symbol].Invested:
                self.Liquidate(symbol) # Sell
                self.Debug(f"Selling {symbol} at {self.Time}, Price: {price}, ARO_DPTL_{self.period}YR: {aro_value}")

            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)
        
        
        # self.portfolioTargets = []
       
# 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/sorted_5_stocks2.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