# region imports
from AlgorithmImports import *
# endregion

class FundamentalBacktester1(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)
        self.SetEndDate(2001, 9, 4)
        self.SetCash(100000) # Set initial cash balance (portfolio size)

        # Universe settings
        self.UniverseSettings.Resolution = Resolution.Daily # Set resolution to daily

        # portolio constuction model
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: None)) # Use equal weighted portfolio construction module
        self.Settings.RebalancePortfolioOnInsightChanges = False # means that the portfoli constuction module above will not update porfio when insights change
        self.Settings.RebalancePortfolioOnSecurityChanges = True # means that the portfoli constuction module above will update porfio when securities change

        # Set universe selection model
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseFilter, self.FineFilter)) # sets the stock universe to whaterver come out of the course & then fine filter fuctions

        # set benchmark
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol # call spy (S&P500) to save in self.spy
        self.SetBenchmark("SPY") # sets spy as the benchmark (aplha determined by outperformance to this)

        # portfolio Rebalance scedule
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY", 120), self.Rebalance)  # call Rebalance function at the start of the month 2hr after market open - base market hours on SPY

        #initalise variables
        self.num_coarse = 500  # number of stocks that will come out of coarse selection filter (used in fuction)
        self.num_fine = 50      # number of stocks that will come out of coarse selection filter (used in fuction)
        self.day_counter = 0
        self.new_symbols = []
        self.filtered_stocks = []
        self.PV = []
        self.nextday = False
        self.monthly_rebalance = False
        self.rebalanced = False

    
    # coarse universe filter fuction (runs every day at midnight)
    #filters stock universe, returns 500 sybmols with the largest market cap, that have available fundamdal data and dollar volume > $5m
    def CoarseFilter(self, coarse):
        if self.nextday:
            return Universe.Unchanged

        # If the rebalance flag is not set, do not change the universe
        if not self.monthly_rebalance:
            self.rebalanced = False
            return Universe.Unchanged 

        self.rebalanced = True
        self.monthly_rebalance = False

        filtered_coarse = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 0 and x.DollarVolume > 5000000],key=lambda x: x.market_cap, reverse=True)   # Filter out all stocks that have no fundamental data or DollarVolume <= $5M, sort by market cap - large to small
        
        # send list of filtered coarse stocks to debugger inc number and pe_ratio
        debug_output = ""
        for i, coarse_obj in enumerate(filtered_coarse[:self.num_coarse], start=1):
            symbol_name = coarse_obj.Symbol.Value  # Ensure to get only the symbol string
            market_cap = coarse_obj.market_cap
            debug_output += f"{i}. {symbol_name}: Market Cap = {market_cap}\n"
        self.Debug(f"Course Filtered Symbols:\n{debug_output}  ")  # Log the selected symbols with P/E ratios

        return [x.Symbol for x in filtered_coarse[:self.num_coarse]]  # Return the top 'num_coarse' symbols
    
    # fine universe filter fuction, (runs every day at midnight)
    #filters Course filter output (500 sybmols) returns 50 sybmols with lowest PE Ratio
    def FineFilter(self, fine):
        # Filter and sort stocks by P/E ratio
        self.filtered_stocks = []

        selected_fine = [x for x in fine if x.valuation_ratios.pe_ratio > 0]

        filtered_fine = sorted([x for x in selected_fine], key=lambda x: x.valuation_ratios.pe_ratio)  # Sort by P/E ratio
        
        self.filtered_stocks = [x.Symbol for x in filtered_fine[:self.num_fine]]  # Select the top 50 'num_fine' symbols

        # send list of filtered fine stocks to debugger inc number and pe_ratio
        debug_output = ""
        for i, fine_obj in enumerate(filtered_fine[:self.num_fine], start=1):
            symbol_name = fine_obj.Symbol.Value  # Ensure to get only the symbol string
            pe_ratio = fine_obj.valuation_ratios.pe_ratio
            debug_output += f"{i}. {symbol_name}: PE Ratio = {pe_ratio}\n"

        self.Debug(f"Fine Filtered Symbols\n{debug_output}  ")  # Log the selected symbols with P/E ratios

        return self.filtered_stocks  # Return the selected symbols

    # OnData fuction (runs every day at midnight)
    def OnData(self, data):
        self.day_counter += 1
        # returns out of function every call exept if just rebalenced
        if not self.rebalanced:
            return

        # If the next day = 1 skip till the next bar/day/on_data call (means that we always buy stock 1 day after getting buy signal)
        if not self.nextday:
            self.nextday = True
            self.Debug("skip day "+str(self.time) +"   day "+str(self.day_counter))
            return
        self.Debug("next day "+str(self.time))
        
        self.nextday = False
        
        insights = []

        # Close insights for symbols no longer in self.filtered_stocks
        for symbol in list(self.PV):  # Create a copy of the list to modify it during iteration
            if symbol not in self.filtered_stocks:
                insights.append(Insight.Price(symbol, timedelta(days=7560), InsightDirection.Flat))
                self.PV.remove(symbol)
                self.Debug(f"Close {symbol}")  # log symbol close

        # Add new insights for symbols not in self.PV (open)
        for symbol in self.filtered_stocks:
            if symbol not in self.PV:
                insights.append(Insight.Price(symbol, timedelta(days=7560), InsightDirection.Up))
                self.Debug(f"Open {symbol}")  # log symbol open
                self.PV.append(symbol)

                
        self.EmitInsights(insights)

    
    # this is called according to Schedule.On in Initialize    
    def Rebalance(self):
        self.monthly_rebalance = True