import tkinter as tk
from tkinter import ttk

import matplotlib.figure
import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.axes import Axes
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

from medicine_manager import MedicineManager
from pathology_manager import PathologyManager


class GraphManager:
    """Optimized version - Handle all graph-related operations for the
    application with performance improvements."""

    def __init__(
        self,
        parent_frame: ttk.LabelFrame,
        medicine_manager: MedicineManager,
        pathology_manager: PathologyManager,
    ) -> None:
        self.parent_frame: ttk.LabelFrame = parent_frame
        self.medicine_manager = medicine_manager
        self.pathology_manager = pathology_manager

        # Initialize matplotlib with optimized settings
        self.fig: matplotlib.figure.Figure = plt.figure(figsize=(10, 6), dpi=80)
        self.ax: Axes = self.fig.add_subplot(111)

        # Cache for current data to avoid reprocessing
        self.current_data: pd.DataFrame = pd.DataFrame()
        self._last_plot_hash: str = ""

        # Initialize UI components
        self.toggle_vars: dict[str, tk.IntVar] = {}
        self._setup_ui()
        self._initialize_toggle_vars()
        self._create_chart_toggles()

    def _initialize_toggle_vars(self) -> None:
        """Initialize toggle variables for chart elements with optimization."""
        # Initialize pathology toggles
        for pathology_key in self.pathology_manager.get_pathology_keys():
            self.toggle_vars[pathology_key] = tk.IntVar(value=1)

        # Initialize medicine toggles (unchecked by default)
        for medicine_key in self.medicine_manager.get_medicine_keys():
            self.toggle_vars[medicine_key] = tk.IntVar(value=0)

    def _setup_ui(self) -> None:
        """Set up the UI components with performance optimizations."""
        # Create canvas with optimized settings
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.parent_frame)
        self.canvas.draw_idle()  # Use draw_idle for better performance

        # Pack canvas
        canvas_widget = self.canvas.get_tk_widget()
        canvas_widget.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

        # Create control frame
        self.control_frame = ttk.Frame(self.parent_frame)
        self.control_frame.pack(side=tk.BOTTOM, fill=tk.X, padx=5, pady=2)

    def _create_chart_toggles(self) -> None:
        """Create toggle controls for chart elements with improved layout."""
        # Pathology toggles
        pathology_frame = ttk.LabelFrame(
            self.control_frame, text="Pathologies", padding="5"
        )
        pathology_frame.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=2)

        # Use grid for better layout
        row, col = 0, 0
        for pathology_key in self.pathology_manager.get_pathology_keys():
            pathology = self.pathology_manager.get_pathology(pathology_key)
            if pathology:
                display_name = pathology.display_name
                text = (
                    display_name[:10] + "..."
                    if len(display_name) > 10
                    else display_name
                )
                cb = ttk.Checkbutton(
                    pathology_frame,
                    text=text,
                    variable=self.toggle_vars[pathology_key],
                    command=self._handle_toggle_changed,
                )
                cb.grid(row=row, column=col, sticky="w", padx=2)
                col += 1
                if col > 1:  # 2 columns max
                    col = 0
                    row += 1

        # Medicine toggles
        medicine_frame = ttk.LabelFrame(
            self.control_frame, text="Medicines", padding="5"
        )
        medicine_frame.pack(side=tk.RIGHT, fill=tk.X, expand=True, padx=2)

        # Use grid for medicines too
        row, col = 0, 0
        for medicine_key in self.medicine_manager.get_medicine_keys():
            medicine = self.medicine_manager.get_medicine(medicine_key)
            if medicine:
                med_name = medicine.display_name
                text = med_name[:10] + "..." if len(med_name) > 10 else med_name
                cb = ttk.Checkbutton(
                    medicine_frame,
                    text=text,
                    variable=self.toggle_vars[medicine_key],
                    command=self._handle_toggle_changed,
                )
                cb.grid(row=row, column=col, sticky="w", padx=2)
                col += 1
                if col > 2:  # 3 columns max for medicines
                    col = 0
                    row += 1

    def _handle_toggle_changed(self) -> None:
        """Handle toggle changes by replotting the graph with optimization."""
        if not self.current_data.empty:
            self._plot_graph_data(self.current_data)

    def update_graph(self, df: pd.DataFrame) -> None:
        """Update the graph with new data using optimization checks."""
        # Create hash of data to avoid unnecessary redraws
        data_hash = str(hash(str(df.values.tobytes()) if not df.empty else "empty"))

        # Only update if data actually changed
        if data_hash != self._last_plot_hash or self.current_data.empty:
            self.current_data = df.copy() if not df.empty else pd.DataFrame()
            self._last_plot_hash = data_hash
            self._plot_graph_data(df)

    def _plot_graph_data(self, df: pd.DataFrame) -> None:
        """Plot the graph data with current toggle settings using optimizations."""
        # Use batch updates to reduce redraws
        with plt.ioff():  # Turn off interactive mode for batch updates
            self.ax.clear()

            if not df.empty:
                # Optimize data processing
                df_processed = self._preprocess_data(df)

                # Track if any series are plotted
                has_plotted_series = self._plot_pathology_data(df_processed)
                medicine_data = self._plot_medicine_data(df_processed)

                if has_plotted_series or medicine_data["has_plotted"]:
                    self._configure_graph_appearance(medicine_data)

        # Single draw call at the end
        self.canvas.draw_idle()

    def _preprocess_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """Preprocess data for plotting with optimizations."""
        df = df.copy()
        # Batch convert dates and sort
        df["date"] = pd.to_datetime(df["date"], cache=True)
        df = df.sort_values(by="date")
        df.set_index(keys="date", inplace=True)
        return df

    def _plot_pathology_data(self, df: pd.DataFrame) -> bool:
        """Plot pathology data series with optimizations."""
        has_plotted_series = False

        # Batch plot pathology data
        pathology_keys = self.pathology_manager.get_pathology_keys()
        active_pathologies = [
            key
            for key in pathology_keys
            if self.toggle_vars[key].get() and key in df.columns
        ]

        for pathology_key in active_pathologies:
            pathology = self.pathology_manager.get_pathology(pathology_key)
            if pathology:
                label = f"{pathology.display_name} ({pathology.scale_info})"
                linestyle = (
                    "dashed" if pathology.scale_orientation == "inverted" else "-"
                )
                self._plot_series(df, pathology_key, label, "o", linestyle)
                has_plotted_series = True

        return has_plotted_series

    def _plot_medicine_data(self, df: pd.DataFrame) -> dict:
        """Plot medicine data with optimizations."""
        result = {"has_plotted": False, "with_data": [], "without_data": []}

        # Get medicine colors and keys in batch
        medicine_colors = self.medicine_manager.get_graph_colors()
        medicines = self.medicine_manager.get_medicine_keys()

        # Pre-calculate daily doses for all medicines to avoid repeated computation
        medicine_doses = {}
        for medicine in medicines:
            dose_column = f"{medicine}_doses"
            if dose_column in df.columns:
                daily_doses = [
                    self._calculate_daily_dose(dose_str) for dose_str in df[dose_column]
                ]
                medicine_doses[medicine] = daily_doses

        # Plot medicines with data
        for medicine in medicines:
            if self.toggle_vars[medicine].get() and medicine in medicine_doses:
                daily_doses = medicine_doses[medicine]

                # Check if there's any data to plot
                if any(dose > 0 for dose in daily_doses):
                    result["with_data"].append(medicine)

                    # Optimize dose scaling and bar plotting
                    scaled_doses = [dose / 10 for dose in daily_doses]

                    # Calculate statistics more efficiently
                    non_zero_doses = [d for d in daily_doses if d > 0]
                    if non_zero_doses:
                        avg_dose = sum(daily_doses) / len(non_zero_doses)
                        label = f"{medicine.capitalize()} (avg: {avg_dose:.1f}mg)"

                        # Single bar plot call
                        self.ax.bar(
                            df.index,
                            scaled_doses,
                            alpha=0.6,
                            color=medicine_colors.get(medicine, "#DDA0DD"),
                            label=label,
                            width=0.6,
                            bottom=-max(scaled_doses) * 1.1 if scaled_doses else -1,
                        )
                        result["has_plotted"] = True
                else:
                    # Medicine is toggled on but has no dose data
                    if self.toggle_vars[medicine].get():
                        result["without_data"].append(medicine)

        return result

    def _configure_graph_appearance(self, medicine_data: dict) -> None:
        """Configure graph appearance with optimizations."""
        # Get legend data in batch
        handles, labels = self.ax.get_legend_handles_labels()

        # Add information about medicines without data if any are toggled on
        if medicine_data["without_data"]:
            med_list = ", ".join(medicine_data["without_data"])
            info_text = f"Tracked (no doses): {med_list}"
            labels.append(info_text)

            # Create dummy handle more efficiently
            from matplotlib.patches import Rectangle

            dummy_handle = Rectangle(
                (0, 0), 1, 1, fc="w", fill=False, edgecolor="none", linewidth=0
            )
            handles.append(dummy_handle)

        # Create legend with optimized settings
        if handles and labels:
            self.ax.legend(
                handles,
                labels,
                loc="upper left",
                bbox_to_anchor=(0, 1),
                ncol=2,
                fontsize="small",
                frameon=True,
                fancybox=True,
                shadow=True,
                framealpha=0.9,
            )

        # Set titles and labels
        self.ax.set_title("Medication Effects Over Time")
        self.ax.set_xlabel("Date")
        self.ax.set_ylabel("Rating (0-10) / Dose (mg)")

        # Optimize y-axis configuration
        current_ylim = self.ax.get_ylim()
        self.ax.set_ylim(bottom=current_ylim[0], top=max(10, current_ylim[1]))

        # Optimize date formatting
        self.fig.autofmt_xdate()

    def _plot_series(
        self,
        df: pd.DataFrame,
        column: str,
        label: str,
        marker: str,
        linestyle: str,
    ) -> None:
        """Helper method to plot a data series with optimizations."""
        # Use more efficient plotting parameters
        self.ax.plot(
            df.index,
            df[column],
            marker=marker,
            linestyle=linestyle,
            label=label,
            markersize=4,  # Smaller markers for better performance
            linewidth=1.5,  # Optimized line width
        )

    def _calculate_daily_dose(self, dose_str: str) -> float:
        """Calculate total daily dose from dose string format with optimizations."""
        if not dose_str or pd.isna(dose_str) or str(dose_str).lower() == "nan":
            return 0.0

        total_dose = 0.0
        # Optimize string processing
        dose_str = str(dose_str).replace("•", "").strip()

        # More efficient splitting and processing
        dose_entries = dose_str.split("|") if "|" in dose_str else [dose_str]

        for entry in dose_entries:
            entry = entry.strip()
            if not entry:
                continue

            try:
                # More efficient dose extraction
                dose_part = entry.split(":")[-1] if ":" in entry else entry

                # Optimized numeric extraction
                dose_value = ""
                for char in dose_part:
                    if char.isdigit() or char == ".":
                        dose_value += char
                    elif dose_value:
                        break

                if dose_value:
                    total_dose += float(dose_value)
            except (ValueError, IndexError):
                continue

        return total_dose

    def close(self) -> None:
        """Clean up resources with proper optimization."""
        try:
            # Clear the plot before closing
            self.ax.clear()
            plt.close(self.fig)
        except Exception:
            pass  # Ignore cleanup errors