# region imports
from AlgorithmImports import *
import pandas as pd
from collections import defaultdict
import io
# endregion


class ComprehensiveEarningsDateComparisonWithTime(QCAlgorithm): # Renamed class slightly

    def Initialize(self):
        # --- Critical Settings ---
        self.object_store_key = "filteredEarnings.csv"
        # Save results under a new key to avoid overwriting previous results
        self.results_save_key = "earnings_comparison_details_with_time.csv"

        # --- Use explicit dates based on your CSV range ---
        self.SetStartDate(2020, 1, 17) # Slightly before first CSV date
        self.SetEndDate(2024, 12, 25)   # Slightly after last CSV date

        self.SetCash(100000)

        # --- Data Structures ---
        # ---> CHANGE 2: Modify data structure to store {date: time_str} mapping <---
        # Structure: { symbol: { datetime.date : "BMO" / "AMC" } }
        self.csv_earnings_by_symbol = defaultdict(dict)
        self.csv_unique_symbols = set()
        self.comparison_summary = {"matches": 0, "mismatches": 0, "qc_events_for_csv_symbols": 0}
        self.detailed_results = []
        self.processed_qc_dates = set()

        # --- CSV Loading ---
        if not self.ObjectStore.ContainsKey(self.object_store_key):
            self.Error(f"'{self.object_store_key}' not found in Object Store.")
            self.Quit()
            return

        try:
            file_path = self.ObjectStore.GetFilePath(self.object_store_key)
            # Load all three columns now
            earnings_data_df = pd.read_csv(file_path, header=None, names=["symbol", "earnings_date", "earnings_time_str"])

            # Parse date correctly, keeping only the date part
            earnings_data_df["date_only"] = pd.to_datetime(earnings_data_df["earnings_date"].str.slice(0, 19), format="%Y-%m-%d %H:%M:%S").dt.date

            # ---> CHANGE 2 (cont.): Populate the modified data structure <---
            processed_count = 0
            for index, row in earnings_data_df.iterrows():
                symbol_str = row["symbol"]
                earnings_dt = row["date_only"]
                time_str_raw = row["earnings_time_str"]

                # Map CSV time string to standardized BMO/AMC
                if "Before market open" in time_str_raw:
                    time_std = "BMO"
                elif "After market close" in time_str_raw:
                    time_std = "AMC"
                else:
                    # Log a warning if unexpected value found in CSV time column
                    self.Debug(f"Warning: Unexpected time string '{time_str_raw}' in CSV for {symbol_str} on {earnings_dt}. Skipping this entry.")
                    continue # Skip rows with unexpected time strings

                # Store date -> time mapping for the symbol
                self.csv_earnings_by_symbol[symbol_str][earnings_dt] = time_std
                self.csv_unique_symbols.add(symbol_str)
                processed_count += 1

            self.Log(f"Loaded and processed {processed_count} earnings records (with valid BMO/AMC) from CSV for {len(self.csv_unique_symbols)} unique symbols.")
            del earnings_data_df # Free up memory

        except Exception as e:
            self.Error(f"Error processing CSV file: {e}")
            self.Quit()
            return
        # --- End CSV Loading ---

        # --- Universe Setup ---
        self.UniverseSettings.Resolution = Resolution.Daily
        # Assuming EODHDUpcomingEarnings provides the ReportTime attribute
        self.AddUniverse(EODHDUpcomingEarnings, self.SelectionFilter)
        # --- End Universe Setup ---

    # Removed get_csv_date_range as per user request

    def SelectionFilter(self, earnings: List[EODHDUpcomingEarnings]) -> List[Symbol]:
        """
        Compares QC earnings dates AND times against CSV data and stores detailed results.
        """
        for earning_event in earnings:
            symbol_str = earning_event.Symbol.Value
            qc_report_date = earning_event.ReportDate.date()
            # ---> CHANGE 3: Get QC Report Time and map it <---
            qc_report_time_enum = earning_event.ReportTime

            if qc_report_time_enum == 0:
                qc_time_std = "BMO"
            elif qc_report_time_enum == 1:
                qc_time_std = "AMC"
            # Add handling for other potential QC times if necessary (e.g., DURING_MARKET)
            # elif qc_report_time_enum == ReportTime.DURING_MARKET_HOURS:
            #     qc_time_std = "DURING" # Decide how to handle this - treat as mismatch?
            else:
                qc_time_std = "UNKNOWN"
                # Only log UNKNOWN once per symbol/date to avoid spamming
                if (symbol_str, qc_report_date, "QC_Time_Unknown") not in self.processed_qc_dates:
                     self.Debug(f"Warning: Unknown/Unhandled QC ReportTime '{qc_report_time_enum}' for {symbol_str} on {qc_report_date}")
                     self.processed_qc_dates.add((symbol_str, qc_report_date, "QC_Time_Unknown"))


            # Use a combined key including time for processing check if needed, but date check is likely sufficient
            event_key = (symbol_str, qc_report_date) # Sticking to date key for processed check

            # Process only relevant symbols & new events
            if symbol_str in self.csv_unique_symbols and event_key not in self.processed_qc_dates:
                self.comparison_summary["qc_events_for_csv_symbols"] += 1
                self.processed_qc_dates.add(event_key)

                # ---> CHANGE 3 (cont.): Perform combined date and time check <---
                csv_date_time_map = self.csv_earnings_by_symbol.get(symbol_str, {})
                csv_expected_time_std = csv_date_time_map.get(qc_report_date, "N/A") # Get expected time if date exists

                match_status = False # Default to mismatch
                if qc_report_date in csv_date_time_map:
                    # Date matches, now check if time matches
                    if csv_expected_time_std == qc_time_std:
                        match_status = True
                    # Else: Date matches but time doesn't -> remains mismatch
                # Else: Date doesn't match -> remains mismatch

                # Store detailed results, including time info
                result_detail = {
                    "Symbol": symbol_str,
                    "QC_Report_Date": qc_report_date,
                    "QC_Report_Time": qc_time_std, # Add QC mapped time
                    "CSV_Expected_Time": csv_expected_time_std, # Add CSV mapped time for this date
                    "Match_Status": match_status
                    # Removed CSV_Dates_For_Symbol_List to keep results simpler,
                    # can be added back if needed for debugging mismatches later
                }
                self.detailed_results.append(result_detail)

                # Update summary counts
                if match_status:
                    self.comparison_summary["matches"] += 1
                else:
                    self.comparison_summary["mismatches"] += 1
                    # Optional: Debug log for mismatches showing times
                    # self.Debug(f"MISMATCH: {symbol_str} {qc_report_date} QC:{qc_time_std} CSV:{csv_expected_time_std}")


        return [] # No need to subscribe to equity data

    def OnEndOfAlgorithm(self):
        """
        Log summary statistics and save detailed results (including time) to Object Store.
        """
        self.Log("--- End of Algorithm: Earnings Date & Time Comparison Summary ---") # Updated title
        total_compared = self.comparison_summary["qc_events_for_csv_symbols"]
        matches = self.comparison_summary["matches"]
        mismatches = self.comparison_summary["mismatches"]

        self.Log(f"Backtest Range: {self.StartDate.date()} to {self.EndDate.date()}")
        self.Log(f"Total QC Earning Events Checked (for symbols in CSV): {total_compared}")
        self.Log(f"Matching Dates & Times Found: {matches}") # Updated label
        self.Log(f"Mismatched Dates or Times Found: {mismatches}") # Updated label

        if total_compared > 0:
            accuracy = (matches / total_compared) * 100
            self.Log(f"Accuracy (Matching Date & Time / Total Checked): {accuracy:.2f}%") # Updated label
        else:
            self.Log("No QC earning events found for symbols listed in the CSV during the backtest period.")

        # --- Save detailed results to Object Store ---
        if self.detailed_results:
            self.Log(f"Preparing to save {len(self.detailed_results)} detailed comparison results...")
            try:
                results_df = pd.DataFrame(self.detailed_results)

                # Convert date objects to strings for reliable CSV storage
                # No complex list columns added this time, simpler conversion.
                results_df['QC_Report_Date'] = results_df['QC_Report_Date'].astype(str)

                # ---> CHANGE 4: Ensure new columns QC_Report_Time, CSV_Expected_Time are saved (automatic with DataFrame) <---
                save_path = self.ObjectStore.Save(self.results_save_key, results_df.to_csv(index=False))

                if save_path:
                    self.Log(f"Detailed comparison results successfully saved to Object Store: {save_path}")
                else:
                     self.Error("Failed to save detailed results to Object Store. Save method returned empty path.")

            except Exception as e:
                self.Error(f"Error saving detailed results to Object Store: {e}")
        else:
            self.Log("No detailed results were generated to save.")
        # --- End Save ---