"""
Solve the quadratic equation ax^2 + bx + c = 0.
"""
import cmath
from fractions import Fraction
import sympy as sp
import gradio as gr
import numpy as np

def solve_quadratic(a: float, b: float, c: float, return_format: str = "string"):
    if a == 0:
        if b == 0:
            result = "Not a valid equation (a and b cannot both be zero)." if c != 0 else "Infinite solutions (0 = 0)"
            return result if return_format == "string" else {"error": result}
        root = -c/b
        if return_format == "string":
            return f"Linear equation: x = {root}"
        else:
            return {"roots": (root, None), "vertex": None}
    vertex_x = -b / (2 * a)
    vertex_y = c - (b**2 / (4 * a))
    vertex = (vertex_x, vertex_y)
    delta = b**2 - 4*a*c
    # Surd form using sympy
    if return_format == "surd":
        x1 = (-b + sp.sqrt(delta)) / (2*a)
        x2 = (-b - sp.sqrt(delta)) / (2*a)
        return {
            "roots": (sp.simplify(x1), sp.simplify(x2)),
            "vertex": vertex,
            "discriminant": delta
        }
    if return_format == "dict":
        # Use numpy.roots for quadratic
        roots = np.roots([a, b, c])
        # Ensure two roots (may be complex)
        if len(roots) == 1:
            roots = (roots[0], None)
        else:
            roots = tuple(roots)
        # Try to convert to Fraction if real
        roots_fmt = []
        for r in roots:
            if r is not None and np.isreal(r):
                roots_fmt.append(Fraction(r.real).limit_denominator())
            else:
                roots_fmt.append(r)
        return {"roots": tuple(roots_fmt), "vertex": vertex}
    if delta > 0:
        x1 = (-b + delta**0.5) / (2*a)
        x2 = (-b - delta**0.5) / (2*a)
        try:
            x1_frac = Fraction(x1).limit_denominator()
            x2_frac = Fraction(x2).limit_denominator()
            return f"Two distinct real roots: x1 = {x1_frac}, x2 = {x2_frac}\nVertex at: {vertex}"
        except:
            return f"Two distinct real roots: x1 = {x1}, x2 = {x2}\nVertex at: {vertex}"
    elif delta == 0:
        x1 = -b / (2*a)
        try:
            x1_frac = Fraction(x1).limit_denominator()
            return f"One real root (repeated): x = {x1_frac}\nVertex at: {vertex}"
        except:
            return f"One real root (repeated): x = {x1}\nVertex at: {vertex}"
    else:
        real_part = -b / (2*a)
        imag_part = (-delta)**0.5 / (2*a)
        return f"Two complex roots: x1 = {real_part} + {imag_part}i, x2 = {real_part} - {imag_part}i\nVertex at: {vertex}"

def quadratic_solver_wrapper(a, b, c, return_format):
    result = solve_quadratic(a, b, c, return_format=return_format)
    if return_format == "dict":
        if "error" in result:
            return result["error"]
        roots = result["roots"]
        vertex = result["vertex"]
        output = ""
        sign_b = "+" if b >= 0 else ""
        sign_c = "+" if c >= 0 else ""
        output += f"Equation: {a}x² {sign_b} {b}x {sign_c} {c} = 0\n\n"
        if roots[1] is None:
            output += f"Root: {roots[0]}\n"
        else:
            output += f"Root 1: {roots[0]}\n"
            output += f"Root 2: {roots[1]}\n"
        if vertex:
            output += f"\nVertex: ({vertex[0]}, {vertex[1]})"
        return output
    else:
        return result

solve_quadratic_interface = gr.Interface(
    fn=quadratic_solver_wrapper,
    inputs=[
        gr.Number(label="a (coefficient of x²)"),
        gr.Number(label="b (coefficient of x)"),
        gr.Number(label="c (constant)"),
        gr.Radio(
            choices=["string", "dict", "surd"],
            value="dict",
            label="Output Format",
            info="'string' for text output, 'dict' for formatted output, 'surd' for exact roots"
        )
    ],
    outputs="text",
    title="Quadratic Equation Solver",
    description="""
Solve ax² + bx + c = 0 and find the vertex. Enter the coefficients for your quadratic equation and select the output format.

Example: For x² - 3x + 2 = 0, enter a=1, b=-3, c=2.

Output format:
- 'string': plain text
- 'dict': formatted output
- 'surd': exact roots (if possible)
""",
    examples=[
        [1, -3, 2, "dict"],
        [2, 4, -6, "string"],
        [1, 2, 1, "surd"]
    ]
)

def plot_quadratic(a, b, c):
    import numpy as np
    import matplotlib.pyplot as plt
    result = solve_quadratic(a, b, c, return_format="dict")
    vertex_x = -b / (2*a) if a != 0 else 0
    vertex_y = c - (b**2 / (4*a)) if a != 0 else 0
    fig, ax = plt.subplots(figsize=(8, 6))
    if a != 0:
        if "roots" in result and result["roots"][0] is not None and result["roots"][1] is not None:
            root1 = result["roots"][0].real if hasattr(result["roots"][0], "real") else float(result["roots"][0])
            root2 = result["roots"][1].real if hasattr(result["roots"][1], "real") else float(result["roots"][1])
            x_min = min(root1, root2, vertex_x) - 2
            x_max = max(root1, root2, vertex_x) + 2
        else:
            x_min = vertex_x - 5
            x_max = vertex_x + 5
        x = np.linspace(x_min, x_max, 1000)
        y = a * (x**2) + b * x + c
        ax.plot(x, y, 'b-', label=f'f(x) = {a}x² + {b}x + {c}')
        ax.plot(vertex_x, vertex_y, 'ro', label=f'Vertex: ({vertex_x:.2f}, {vertex_y:.2f})')
        if "roots" in result:
            roots = result["roots"]
            if roots[0] is not None and (isinstance(roots[0], (int, float)) or (hasattr(roots[0], "imag") and roots[0].imag == 0)):
                root1 = float(roots[0].real if hasattr(roots[0], "real") else roots[0])
                ax.plot(root1, 0, 'go', label=f'Root 1: {root1:.2f}')
            if roots[1] is not None and (isinstance(roots[1], (int, float)) or (hasattr(roots[1], "imag") and roots[1].imag == 0)):
                root2 = float(roots[1].real if hasattr(roots[1], "real") else roots[1])
                ax.plot(root2, 0, 'go', label=f'Root 2: {root2:.2f}')
        ax.axhline(y=0, color='k', linestyle='-', alpha=0.3)
        ax.axvline(x=0, color='k', linestyle='-', alpha=0.3)
        ax.grid(True, alpha=0.3)
        ax.set_xlabel('x')
        ax.set_ylabel('f(x)')
        ax.set_title(f'Graph of f(x) = {a}x² + {b}x + {c}')
        ax.legend()
    else:
        if b != 0:
            x = np.linspace(-5, 5, 100)
            y = b * x + c
            ax.plot(x, y, 'b-', label=f'f(x) = {b}x + {c} (Linear)')
            root = -c/b
            ax.plot(root, 0, 'go', label=f'Root: {root:.2f}')
            ax.axhline(y=0, color='k', linestyle='-', alpha=0.3)
            ax.axvline(x=0, color='k', linestyle='-', alpha=0.3)
            ax.grid(True, alpha=0.3)
            ax.set_xlabel('x')
            ax.set_ylabel('f(x)')
            ax.set_title(f'Graph of f(x) = {b}x + {c} (Linear)')
            ax.legend()
        else:
            x = np.linspace(-5, 5, 100)
            y = c * np.ones_like(x)
            ax.plot(x, y, 'b-', label=f'f(x) = {c} (Constant)')
            ax.axhline(y=0, color='k', linestyle='-', alpha=0.3)
            ax.axvline(x=0, color='k', linestyle='-', alpha=0.3)
            ax.grid(True, alpha=0.3)
            ax.set_xlabel('x')
            ax.set_ylabel('f(x)')
            ax.set_title(f'Graph of f(x) = {c} (Constant)')
            ax.legend()
    return fig

quadratic_visualizer_interface = gr.Interface(
    fn=plot_quadratic,
    inputs=[
        gr.Number(label="a (coefficient of x²)", value=1),
        gr.Number(label="b (coefficient of x)", value=0),
        gr.Number(label="c (constant)", value=0)
    ],
    outputs=gr.Plot(),
    title="Quadratic Function Visualizer",
    description="""
Visualize the graph of a quadratic function f(x) = ax² + bx + c, including its vertex and real roots (if any).

Example: For f(x) = x² - 4x + 3, enter a=1, b=-4, c=3.
""",
    examples=[
        [1, -4, 3],
        [2, 0, -8],
        [1, 2, 1]
    ]
)