import tkinter as tk
from tkinter import ttk, messagebox
import numpy as np
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

# Import NavigationToolbar2Tk with version compatibility
try:
    # For matplotlib >= 3.4
    from matplotlib.backends.backend_tkagg import NavigationToolbar2Tk  # type: ignore
except ImportError:
    try:
        # For older matplotlib versions or alternative location
        from matplotlib.backends._backend_tk import NavigationToolbar2Tk  # type: ignore
    except ImportError:
        # Try another possible location
        from matplotlib.backends.backend_tk import NavigationToolbar2Tk  # type: ignore


class QuadraticPlotter:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("Quadratic Equation Plotter")
        self.root.geometry("900x700")
        
        # Variables for coefficients
        self.a_var = tk.StringVar(value="1")
        self.b_var = tk.StringVar(value="0")
        self.c_var = tk.StringVar(value="0")
        
        # Variables for x-range
        self.x_min_var = tk.StringVar(value="-10")
        self.x_max_var = tk.StringVar(value="10")
        
        # Create the main window content
        self.create_widgets()
        
        # Start the application loop
        self.root.mainloop()

    def create_widgets(self):
        # Create main frames
        control_frame = tk.Frame(self.root, padx=10, pady=10)
        control_frame.pack(side=tk.TOP, fill=tk.X)
        
        plot_frame = tk.Frame(self.root)
        plot_frame.pack(side=tk.BOTTOM, fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # Title
        title_label = tk.Label(control_frame, text="Quadratic Equation Plotter: y = ax² + bx + c", 
                               font=("Arial", 14, "bold"))
        title_label.grid(row=0, column=0, columnspan=6, pady=(0, 15))
        
        # Coefficients input
        coeff_frame = tk.LabelFrame(control_frame, text="Coefficients", padx=10, pady=10)
        coeff_frame.grid(row=1, column=0, columnspan=6, sticky="ew", pady=(0, 10))
        
        # a coefficient
        tk.Label(coeff_frame, text="a =").grid(row=0, column=0, padx=(0, 5))
        a_entry = tk.Entry(coeff_frame, textvariable=self.a_var, width=10)
        a_entry.grid(row=0, column=1, padx=(0, 15))
        
        # b coefficient
        tk.Label(coeff_frame, text="b =").grid(row=0, column=2, padx=(0, 5))
        b_entry = tk.Entry(coeff_frame, textvariable=self.b_var, width=10)
        b_entry.grid(row=0, column=3, padx=(0, 15))
        
        # c coefficient
        tk.Label(coeff_frame, text="c =").grid(row=0, column=4, padx=(0, 5))
        c_entry = tk.Entry(coeff_frame, textvariable=self.c_var, width=10)
        c_entry.grid(row=0, column=5, padx=(0, 15))
        
        # X-range input
        range_frame = tk.LabelFrame(control_frame, text="X-axis Range", padx=10, pady=10)
        range_frame.grid(row=2, column=0, columnspan=6, sticky="ew", pady=(0, 10))
        
        tk.Label(range_frame, text="From:").grid(row=0, column=0, padx=(0, 5))
        xmin_entry = tk.Entry(range_frame, textvariable=self.x_min_var, width=10)
        xmin_entry.grid(row=0, column=1, padx=(0, 15))
        
        tk.Label(range_frame, text="To:").grid(row=0, column=2, padx=(0, 5))
        xmax_entry = tk.Entry(range_frame, textvariable=self.x_max_var, width=10)
        xmax_entry.grid(row=0, column=3, padx=(0, 15))
        
        # Buttons
        button_frame = tk.Frame(control_frame)
        button_frame.grid(row=3, column=0, columnspan=6, pady=(10, 0))
        
        plot_button = tk.Button(button_frame, text="Plot Quadratic", command=self.plot_quadratic,
                                bg="lightblue", font=("Arial", 10, "bold"), padx=20, pady=5)
        plot_button.pack(side=tk.LEFT, padx=(0, 10))
        
        save_button = tk.Button(button_frame, text="Save Plot", command=self.save_plot,
                                bg="lightgreen", font=("Arial", 10, "bold"), padx=20, pady=5)
        save_button.pack(side=tk.LEFT, padx=(0, 10))
        
        clear_button = tk.Button(button_frame, text="Clear Plot", command=self.clear_plot,
                                bg="lightcoral", font=("Arial", 10, "bold"), padx=20, pady=5)
        clear_button.pack(side=tk.LEFT)
        
        # Matplotlib figure and canvas
        self.fig = Figure(figsize=(8, 6), dpi=100)
        self.ax = self.fig.add_subplot(111)
        self.setup_plot()
        
        self.canvas = FigureCanvasTkAgg(self.fig, master=plot_frame)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        
        # Add matplotlib navigation toolbar
        self.toolbar = NavigationToolbar2Tk(self.canvas, plot_frame)
        self.toolbar.update()
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

    def setup_plot(self):
        """Initialize plot with labels and grid"""
        self.ax.set_xlabel('x')
        self.ax.set_ylabel('y')
        self.ax.set_title('Quadratic Function: y = ax² + bx + c')
        self.ax.grid(True, linestyle='--', alpha=0.7)
        self.ax.axhline(y=0, color='k', linestyle='-', linewidth=0.5)
        self.ax.axvline(x=0, color='k', linestyle='-', linewidth=0.5)

    def plot_quadratic(self):
        """Plot the quadratic equation based on user input"""
        try:
            # Get coefficients
            a = float(self.a_var.get())
            b = float(self.b_var.get())
            c = float(self.c_var.get())
            
            # Get x-range
            x_min = float(self.x_min_var.get())
            x_max = float(self.x_max_var.get())
            
            if x_min >= x_max:
                messagebox.showerror("Input Error", "X-min must be less than X-max")
                return
            
            # Generate x values
            x = np.linspace(x_min, x_max, 500)
            y = a * x**2 + b * x + c
            
            # Clear previous plot
            self.ax.clear()
            self.setup_plot()
            
            # Determine equation type for title
            if a != 0:
                eq_type = "Quadratic Function"
            elif b != 0:
                eq_type = "Linear Function"
            else:
                eq_type = "Constant Function"
            
            # Build equation string with only non-zero terms, formatted nicely
            parts = []
            if a != 0:
                if a == 1:
                    parts.append("x²")
                elif a == -1:
                    parts.append("-x²")
                else:
                    parts.append(f"{a}x²")
            if b != 0:
                if b == 1:
                    parts.append("+ x" if parts else "x")
                elif b == -1:
                    parts.append("- x" if parts else "-x")
                else:
                    sign = "+" if b > 0 else "-"
                    parts.append(f"{sign} {abs(b)}x")
            if c != 0:
                sign = "+" if c > 0 else "-"
                parts.append(f"{sign} {abs(c)}")
            
            if not parts:
                eq_str = "y = 0"
            else:
                # First term without leading plus
                first = parts[0]
                if first.startswith("+ "):
                    first = first[2:]
                elif first.startswith("- "):
                    first = f"-{first[2:]}"
                eq_str = "y = " + first + " " + " ".join(parts[1:])
            
            self.ax.set_title(f"{eq_type}: {eq_str}")
            
            # Plot the quadratic
            self.ax.plot(x, y, 'b-', linewidth=2, label=eq_str)
            
            # Calculate and plot vertex
            if a != 0:
                vertex_x = -b / (2 * a)
                vertex_y = a * vertex_x**2 + b * vertex_x + c
                self.ax.plot(vertex_x, vertex_y, 'ro', markersize=8, label=f'Vertex ({vertex_x:.2f}, {vertex_y:.2f})')
            
            # Calculate and plot roots if they exist (real roots)
            discriminant = b**2 - 4*a*c
            
            # Handle quadratic case (a ≠ 0)
            if a != 0:
                if discriminant >= 0:
                    root1 = (-b + np.sqrt(discriminant)) / (2 * a)
                    root2 = (-b - np.sqrt(discriminant)) / (2 * a)
                    # For discriminant == 0, both roots are the same
                    if abs(root1 - root2) < 1e-10:  # Essentially equal
                        self.ax.plot(root1, 0, 'go', markersize=8, label=f'Double Root ({root1:.2f}, 0)')
                    else:
                        self.ax.plot(root1, 0, 'go', markersize=8, label=f'Root1 ({root1:.2f}, 0)')
                        self.ax.plot(root2, 0, 'go', markersize=8, label=f'Root2 ({root2:.2f}, 0)')
                else:
                    # Complex roots - show message in legend
                    self.ax.text(0.02, 0.02, 'Complex Roots', transform=self.ax.transAxes,
                                fontsize=10, color='red', bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.5))
            # Handle linear case (a = 0, b ≠ 0)
            elif b != 0:
                root = -c / b
                self.ax.plot(root, 0, 'go', markersize=8, label=f'Root ({root:.2f}, 0)')
            # Handle constant case (a = 0, b = 0)
            else:
                if c == 0:
                    self.ax.text(0.02, 0.02, 'Infinite Roots (y = 0)', transform=self.ax.transAxes,
                                fontsize=10, color='blue', bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.5))
                else:
                    self.ax.text(0.02, 0.02, 'No Roots (horizontal line)', transform=self.ax.transAxes,
                                fontsize=10, color='red', bbox=dict(boxstyle='round', facecolor='lightcoral', alpha=0.5))
            
            # Add legend
            self.ax.legend(loc='best')
            
            # Refresh canvas
            self.canvas.draw()
            
        except ValueError:
            messagebox.showerror("Input Error", "Please enter valid numbers for all coefficients and range")

    def clear_plot(self):
        """Clear the plot"""
        self.ax.clear()
        self.setup_plot()
        self.canvas.draw()

    def save_plot(self):
        """Save the current plot to a file"""
        try:
            from tkinter import filedialog
            import os
            
            filetypes = [
                ('PNG files', '*.png'),
                ('JPEG files', '*.jpg'),
                ('PDF files', '*.pdf'),
                ('SVG files', '*.svg'),
                ('All files', '*.*')
            ]
            
            filename = filedialog.asksaveasfilename(
                defaultextension=".png",
                filetypes=filetypes,
                title="Save Plot As"
            )
            
            if filename:
                self.fig.savefig(filename, dpi=300, bbox_inches='tight')
                messagebox.showinfo("Success", f"Plot saved successfully as:\n{os.path.basename(filename)}")
                
        except Exception as e:
            messagebox.showerror("Save Error", f"Could not save plot: {str(e)}")


if __name__ == "__main__":
    app = QuadraticPlotter()