import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import sympy as sp

x, y = sp.symbols('x y')

# Generate polynomial template
def generate_polynomial_template(degree):
    terms = [f"a{i}*x^{degree - i}" for i in range(degree)] + [f"a{degree}"]
    return " + ".join(terms) + " = 0"

# Solve and plot polynomial
def solve_polynomial(degree, coeff_string):
    try:
        coeffs = [sp.sympify(s) for s in coeff_string.strip().split()]
        if len(coeffs) != degree + 1:
            return f"⚠️ Please enter exactly {degree + 1} coefficients.", None, None

        poly = sum([coeffs[i] * x**(degree - i) for i in range(degree + 1)])
        simplified = sp.simplify(poly)

        # Factor step-by-step
        factored_steps = []
        current_expr = simplified
        while True:
            factored = sp.factor(current_expr)
            if factored == current_expr:
                break
            factored_steps.append(factored)
            current_expr = factored

        roots = sp.solve(sp.Eq(simplified, 0), x)
        root_display = []
        for i, r in enumerate(roots):
            r_simplified = sp.nsimplify(r, rational=True)
            root_display.append(f"r_{{{i+1}}} = {sp.latex(r_simplified)}")

        steps_output = f"### 🧐 Polynomial Expression\n\n$$ {sp.latex(poly)} = 0 $$\n\n"
        steps_output += f"### ✏️ Simplified\n\n$$ {sp.latex(simplified)} = 0 $$\n\n"

        if factored_steps:
            steps_output += f"### 🪜 Step-by-Step Factorization\n\n"
            for i, step in enumerate(factored_steps, 1):
                steps_output += f"**Step {i}:** $$ {sp.latex(step)} = 0 $$\n\n"
        else:
            steps_output += f"### 🤷 No further factorization possible\n\n"

        steps_output += "### 🥮 Roots\n\n$$ " + " \\quad ".join(root_display) + " $$"

        f_lambdified = sp.lambdify(x, simplified, modules=["numpy"])
        x_vals = np.linspace(-10, 10, 400)
        y_vals = f_lambdified(x_vals)

        fig, ax = plt.subplots(figsize=(6, 4))
        ax.plot(x_vals, y_vals, label="Polynomial")
        ax.axhline(0, color='black', linewidth=0.5)
        ax.axvline(0, color='black', linewidth=0.5)
        ax.grid(True)
        ax.set_title("📈 Graph of the Polynomial")
        ax.set_xlabel("x")
        ax.set_ylabel("f(x)")

        real_roots = [sp.N(r.evalf()) for r in roots if sp.im(r) == 0]
        for r in real_roots:
            ax.plot([float(r)], [0], 'ro', label="Real Root")

        ax.legend()

        return steps_output, fig, ""

    except Exception as e:
        return f"❌ Error: {e}", None, ""

# Solve linear system
def solve_linear_system(eq1_str, eq2_str):
    try:
        eq1 = sp.sympify(eq1_str)
        eq2 = sp.sympify(eq2_str)

        sol = sp.solve((eq1, eq2), (x, y), dict=True)

        steps = "### 🔍 Solving System\n\n"
        steps += f"**Equation 1:** $$ {sp.latex(eq1)} $$\n\n"
        steps += f"**Equation 2:** $$ {sp.latex(eq2)} $$\n\n"

        if sol:
            sol = sol[0]
            steps += f"**Solution:** $$ x = {sp.latex(sol[x])}, \ y = {sp.latex(sol[y])} $$\n\n"
        else:
            steps += "**No unique solution or inconsistent system**\n"

        x_vals = np.linspace(-10, 10, 400)
        f1 = sp.solve(eq1, y)
        f2 = sp.solve(eq2, y)

        fig, ax = plt.subplots(figsize=(6, 4))
        if f1 and f2:
            y1 = sp.lambdify(x, f1[0], modules=['numpy'])(x_vals)
            y2 = sp.lambdify(x, f2[0], modules=['numpy'])(x_vals)
            ax.plot(x_vals, y1, label='Equation 1')
            ax.plot(x_vals, y2, label='Equation 2')

            if sol:
                px = float(sp.N(sol[x]))
                py = float(sp.N(sol[y]))
                ax.plot(px, py, 'ro')
                ax.annotate(f"({px:.2f}, {py:.2f})", (px, py), textcoords="offset points", xytext=(10, 5), ha='center', color='red')

        ax.axhline(0, color='black', linewidth=0.5)
        ax.axvline(0, color='black', linewidth=0.5)
        ax.set_title("📉 Graph of the Linear System")
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.grid(True)
        ax.legend()

        return steps, fig

    except Exception as e:
        return f"❌ Error: {e}", None

# UI
with gr.Blocks() as demo:
    gr.Markdown("## 🔢 Polynomial Solver with Step-by-Step Factorization and Graph")

    with gr.Row():
        degree_slider = gr.Slider(1, 8, value=3, step=1, label="Degree of Polynomial")
        template_display = gr.Textbox(label="Polynomial Template (Fill in Coefficients)", interactive=False)

    coeff_input = gr.Textbox(label="Enter Coefficients (space-separated, supports pi, e, sqrt(2), I)", placeholder="e.g. 1 -3 sqrt(2) -pi")

    steps_md = gr.Markdown()
    plot_output = gr.Plot()
    error_box = gr.Textbox(visible=False)

    with gr.Row():
        solve_button = gr.Button("Plot Polynomial", variant="primary")

    degree_slider.change(fn=generate_polynomial_template, inputs=degree_slider, outputs=template_display)
    solve_button.click(fn=solve_polynomial, inputs=[degree_slider, coeff_input], outputs=[steps_md, plot_output, error_box])

    gr.Markdown("## 📐 Solve Linear System (2 Equations, 2 Variables)")

    eq1_input = gr.Textbox(label="Equation 1 (in x and y)", placeholder="e.g. 2*x + 3*y - 6")
    eq2_input = gr.Textbox(label="Equation 2 (in x and y)", placeholder="e.g. -x + y - 2")
    sys_steps = gr.Markdown()
    sys_plot = gr.Plot()

    with gr.Row():
        solve_sys_button = gr.Button("Solve Linear System", variant="primary")

    solve_sys_button.click(fn=solve_linear_system, inputs=[eq1_input, eq2_input], outputs=[sys_steps, sys_plot])

    demo.launch()