from AlgorithmImports import *

class EarningsATRUniverseAlgorithm(QCAlgorithm):

    def Initialize(self):
        # sample start date
        self.set_start_date(2024, 1, 1)
        # sample end date
        self.set_end_date(2024, 10, 31)

        self.set_cash(100000)

        # get 2024 earnings calendar from object store, can store custom data in object store
        self.df_earnings = self.get_df_from_store('sylvaint/earnings_2024.csv')
        
        # List to store filtered symbols and dates for earnings and ATR checks
        self.filtered_symbols = []
        # set resolution to Daily
        self.universe_settings.resolution = Resolution.DAILY
        
        # add stock universe with coarse filter function
        self.add_universe(self.CoarseSelectionFunction)

        # Dictionary to track symbols' ATR and earnings previous close
        self.symbol_data = {}
        self.prev_close = {}

    # check symbol has earnings for date
    def is_stock_earnings(self, date_string, symbol):
        try:
            stock_array = (
                self.df_earnings
                .query('date==@date_string')
                .assign(symbols=lambda df_: df_['symbols'].apply(lambda x: x.split(',')))
                .symbols
                .to_list()[0]
            )
        except:
            return False
            
        return symbol in stock_array

    # read earnings dataframe from store
    def get_df_from_store(self, key):
        file_path = self.object_store.get_file_path(key)
        return pd.read_csv(file_path)

    # convert self.time to date string
    def time_to_date_string(self, t):
        return t.strftime('%Y-%m-%d')

    # this is the coarse filter to our stock universe, it is run for every date in our sample range
    def CoarseSelectionFunction(self, coarse):
        # universe selection occurs at midnight so previous day data
        # print the current algorithm date
        self.Debug(f"Coarse selection date: {self.time}")

        selected = []
        
        for stock in coarse:
            # Filter by price > 10 and volume > 500,000 and earnings today
            if (stock.Price > 10 and 
                stock.Volume > 500000 and
                # this checks if stock reported earnings yesterday
                self.is_stock_earnings(self.time_to_date_string(self.time - timedelta(days=1)), stock.symbol.value)):

                # store previous close
                self.prev_close[stock.Symbol.value] = stock.adjusted_price
                selected.append(stock.Symbol)
        # self.debug(f"{len(selected)} stocks from coarse")
        
        return selected

    # this is called every day at midnight and registers added and removed securities according to the coarse filter
    def OnSecuritiesChanged(self, changes):
        # Subscribe to selected securities from universe filter
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            # Initialize ATR and fundamental data tracking
            atr = self.atr(symbol, 14, Resolution.DAILY)
            # Populate ATR with 14 days of historical data
            indicator_history = self.indicator_history(atr, symbol, 14, Resolution.DAILY)
            if indicator_history.Current:
                atr_value = indicator_history.Current[0].value
            else:
                atr_value = None
            self.symbol_data[symbol] = {
                'ATR': atr_value,
            }

        # remove atr until next occurence
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            del self.symbol_data[symbol]
    
    # this is where data comes in, in this case daily data
    def on_data(self, data):
        # Process earnings and ATR conditions for filtered stocks
        for symbol, sym_data in self.symbol_data.items():
            if symbol in data.Bars and sym_data['ATR'] is not None:
                # Get the current open, close, and ATR values
                open_price = data.Bars[symbol].Open
                prev_close = self.prev_close[symbol.value]
                atr_value = sym_data['ATR']
                
                # Check if the stock had earnings yesterday and opened >2 ATRs above the close
                if open_price > (prev_close + 2 * atr_value):
                    self.filtered_symbols.append({
                        'Symbol': symbol.Value,
                        'Date': self.Time,
                        'ATR': atr_value,
                        'PrevClose': prev_close,
                        'Open': open_price
                    })
                
    # this is called at the end of the algo
    def on_end_of_algorithm(self):
        # Display filtered results at the end of the backtest
        for result in self.filtered_symbols:
            # self.log(f"{result['Date']} - {result['Symbol']} - ATR: {result['ATR']} - prev close {result['PrevClose']} - Open {result['Open']} - opened more than 2 ATRs above the previous close.")
            self.log(f"{result['Date']},{result['Symbol']},{result['ATR']},{result['PrevClose']},{result['Open']}")