app.py

import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from matplotlib.patches import Circle
import io
import base64
from PIL import Image
import warnings
warnings.filterwarnings('ignore')

class MathVisualizer:
    def __init__(self):
        self.x = sp.Symbol('x')
        self.y = sp.Symbol('y')
        self.t = sp.Symbol('t')

    def safe_eval(self, expression, variables):
        """Safely evaluate mathematical expressions"""
        try:
            # Convert string to sympy expression
            expr = sp.sympify(expression)
            # Convert to numpy function for evaluation
            func = sp.lambdify(variables, expr, 'numpy')
            return func
        except Exception as e:
            raise ValueError(f"Invalid expression: {str(e)}")

    def plot_2d_function(self, equation, x_range, y_range, color, style, grid, title):
        """Plot 2D function y = f(x)"""
        try:
            func = self.safe_eval(equation, self.x)
            x_vals = np.linspace(x_range[0], x_range[1], 1000)
            y_vals = func(x_vals)
            
            # Handle complex numbers and infinities
            y_vals = np.real(y_vals)
            y_vals = np.where(np.abs(y_vals) > 1e10, np.nan, y_vals)
            
            plt.figure(figsize=(10, 8))
            plt.plot(x_vals, y_vals, color=color, linewidth=2, linestyle=style)
            plt.xlim(x_range)
            plt.ylim(y_range)
            plt.xlabel('x', fontsize=12)
            plt.ylabel('y', fontsize=12)
            plt.title(title or f'y = {equation}', fontsize=14)
            plt.grid(grid, alpha=0.3)
            
            return plt.gcf()
            
        except Exception as e:
            plt.figure(figsize=(10, 8))
            plt.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', 
                    transform=plt.gca().transAxes, fontsize=12, color='red')
            plt.title('Error in equation')
            return plt.gcf()

    def plot_parametric(self, x_equation, y_equation, t_range, color, style, grid, title):
        """Plot parametric equations x = f(t), y = g(t)"""
        try:
            x_func = self.safe_eval(x_equation, self.t)
            y_func = self.safe_eval(y_equation, self.t)
            
            t_vals = np.linspace(t_range[0], t_range[1], 1000)
            x_vals = x_func(t_vals)
            y_vals = y_func(t_vals)
            
            # Handle complex numbers
            x_vals = np.real(x_vals)
            y_vals = np.real(y_vals)
            
            plt.figure(figsize=(10, 8))
            plt.plot(x_vals, y_vals, color=color, linewidth=2, linestyle=style)
            plt.xlabel('x', fontsize=12)
            plt.ylabel('y', fontsize=12)
            plt.title(title or f'x = {x_equation}, y = {y_equation}', fontsize=14)
            plt.grid(grid, alpha=0.3)
            plt.axis('equal')
            
            return plt.gcf()
            
        except Exception as e:
            plt.figure(figsize=(10, 8))
            plt.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', 
                    transform=plt.gca().transAxes, fontsize=12, color='red')
            plt.title('Error in parametric equations')
            return plt.gcf()

    def plot_polar(self, equation, theta_range, color, style, grid, title):
        """Plot polar equations r = f(θ)"""
        try:
            theta = sp.Symbol('theta')
            func = self.safe_eval(equation.replace('θ', 'theta').replace('theta', 'theta'), theta)
            
            theta_vals = np.linspace(theta_range[0], theta_range[1], 1000)
            r_vals = func(theta_vals)
            r_vals = np.real(r_vals)
            
            plt.figure(figsize=(10, 8))
            ax = plt.subplot(111, projection='polar')
            ax.plot(theta_vals, r_vals, color=color, linewidth=2, linestyle=style)
            ax.set_title(title or f'r = {equation}', fontsize=14, pad=20)
            ax.grid(grid, alpha=0.3)
            
            return plt.gcf()
            
        except Exception as e:
            plt.figure(figsize=(10, 8))
            plt.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', 
                    transform=plt.gca().transAxes, fontsize=12, color='red')
            plt.title('Error in polar equation')
            return plt.gcf()

    def plot_implicit(self, equation, x_range, y_range, color, grid, title):
        """Plot implicit equations F(x,y) = 0"""
        try:
            # Parse equation (assume it equals 0)
            expr = sp.sympify(equation)
            
            x_vals = np.linspace(x_range[0], x_range[1], 400)
            y_vals = np.linspace(y_range[0], y_range[1], 400)
            X, Y = np.meshgrid(x_vals, y_vals)
            
            # Convert to numpy function
            func = sp.lambdify([self.x, self.y], expr, 'numpy')
            Z = func(X, Y)
            Z = np.real(Z)
            
            plt.figure(figsize=(10, 8))
            plt.contour(X, Y, Z, levels=[0], colors=[color], linewidths=2)
            plt.xlim(x_range)
            plt.ylim(y_range)
            plt.xlabel('x', fontsize=12)
            plt.ylabel('y', fontsize=12)
            plt.title(title or f'{equation} = 0', fontsize=14)
            plt.grid(grid, alpha=0.3)
            
            return plt.gcf()
            
        except Exception as e:
            plt.figure(figsize=(10, 8))
            plt.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', 
                    transform=plt.gca().transAxes, fontsize=12, color='red')
            plt.title('Error in implicit equation')
            return plt.gcf()

    def plot_3d_surface(self, equation, x_range, y_range, color_scheme, title):
        """Plot 3D surface z = f(x,y)"""
        try:
            func = self.safe_eval(equation, [self.x, self.y])
            
            x_vals = np.linspace(x_range[0], x_range[1], 50)
            y_vals = np.linspace(y_range[0], y_range[1], 50)
            X, Y = np.meshgrid(x_vals, y_vals)
            Z = func(X, Y)
            Z = np.real(Z)
            
            fig = plt.figure(figsize=(12, 10))
            ax = fig.add_subplot(111, projection='3d')
            surf = ax.plot_surface(X, Y, Z, cmap=color_scheme, alpha=0.8)
            ax.set_xlabel('x', fontsize=12)
            ax.set_ylabel('y', fontsize=12)
            ax.set_zlabel('z', fontsize=12)
            ax.set_title(title or f'z = {equation}', fontsize=14)
            fig.colorbar(surf, shrink=0.5, aspect=5)
            
            return fig
            
        except Exception as e:
            fig = plt.figure(figsize=(12, 10))
            plt.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', 
                    transform=plt.gca().transAxes, fontsize=12, color='red')
            plt.title('Error in 3D equation')
            return fig

# Initialize visualizer
visualizer = MathVisualizer()

def generate_plot(plot_type, equation, x_equation, y_equation,
                  x_min, x_max, y_min, y_max, t_min, t_max, theta_min, theta_max,
                  color, line_style, color_scheme, show_grid, custom_title):
    
    plt.close('all')  # Close any existing plots

    try:
        if plot_type == "2D Function (y = f(x))":
            fig = visualizer.plot_2d_function(
                equation, (x_min, x_max), (y_min, y_max), 
                color, line_style, show_grid, custom_title
            )
        
        elif plot_type == "Parametric (x = f(t), y = g(t))":
            fig = visualizer.plot_parametric(
                x_equation, y_equation, (t_min, t_max), 
                color, line_style, show_grid, custom_title
            )
        
        elif plot_type == "Polar (r = f(θ))":
            fig = visualizer.plot_polar(
                equation, (theta_min, theta_max), 
                color, line_style, show_grid, custom_title
            )
        
        elif plot_type == "Implicit (F(x,y) = 0)":
            fig = visualizer.plot_implicit(
                equation, (x_min, x_max), (y_min, y_max), 
                color, show_grid, custom_title
            )
        
        elif plot_type == "3D Surface (z = f(x,y))":
            fig = visualizer.plot_3d_surface(
                equation, (x_min, x_max), (y_min, y_max), 
                color_scheme, custom_title
            )
        
        return fig
        
    except Exception as e:
        plt.figure(figsize=(10, 8))
        plt.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', 
                transform=plt.gca().transAxes, fontsize=12, color='red')
        plt.title('Error generating plot')
        return plt.gcf()

# Example equations for different types
examples = {
    "2D Function (y = f(x))": [
        "sin(x)",
        "x**2 + 2*x + 1",
        "exp(-x**2)",
        "tan(x)",
        "log(abs(x))"
    ],
    "Parametric (x = f(t), y = g(t))": [
        ("cos(t)", "sin(t)"),  # Circle
        ("t*cos(t)", "t*sin(t)"),  # Spiral
        ("cos(3*t)", "sin(2*t)"),  # Lissajous
    ],
    "Polar (r = f(θ))": [
        "1 + cos(theta)",  # Cardioid
        "sin(4*theta)",  # Rose
        "theta",  # Spiral
    ],
    "Implicit (F(x,y) = 0)": [
        "x**2 + y**2 - 1",  # Circle
        "(x**2 + y**2)**2 - 2*(x**2 - y**2)",  # Lemniscate
        "x**3 + y**3 - 3*x*y",  # Folium of Descartes
    ],
    "3D Surface (z = f(x,y))": [
        "sin(sqrt(x**2 + y**2))",
        "x**2 - y**2",
        "exp(-(x**2 + y**2))",
    ]
}

def load_example(plot_type, example_idx):
    if plot_type in examples:
        example_list = examples[plot_type]
        if 0 <= example_idx < len(example_list):
            example = example_list[example_idx]
            if plot_type == "Parametric (x = f(t), y = g(t))":
                return example[0], example[1], "", ""
            else:
                return example, "", "", ""
    return "", "", "", ""

# Create Gradio interface
with gr.Blocks(title="Mathematical Equation Visualizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# 📊 Mathematical Equation Visualizer

Generate beautiful visualizations of mathematical equations with various plot types and customization options.

**Supported functions:** sin, cos, tan, exp, log, sqrt, abs, and basic arithmetic (+, -, *, /, **)
""")
    
    with gr.Row():
        with gr.Column(scale=1):
            plot_type = gr.Dropdown(
                choices=[
                    "2D Function (y = f(x))",
                    "Parametric (x = f(t), y = g(t))",
                    "Polar (r = f(θ))",
                    "Implicit (F(x,y) = 0)",
                    "3D Surface (z = f(x,y))"
                ],
                value="2D Function (y = f(x))",
                label="Plot Type"
            )
            
            with gr.Group():
                equation = gr.Textbox(
                    value="sin(x)",
                    label="Equation",
                    placeholder="e.g., sin(x), x**2 + 1, etc."
                )
                
                with gr.Row(visible=False) as parametric_inputs:
                    x_equation = gr.Textbox(
                        label="x = f(t)",
                        placeholder="e.g., cos(t)"
                    )
                    y_equation = gr.Textbox(
                        label="y = g(t)",
                        placeholder="e.g., sin(t)"
                    )
            
            with gr.Group():
                gr.Markdown("### Range Settings")
                with gr.Row():
                    x_min = gr.Number(value=-10, label="x min")
                    x_max = gr.Number(value=10, label="x max")
                with gr.Row():
                    y_min = gr.Number(value=-10, label="y min")
                    y_max = gr.Number(value=10, label="y max")
                with gr.Row():
                    t_min = gr.Number(value=0, label="t min")
                    t_max = gr.Number(value=6.28, label="t max")
                with gr.Row():
                    theta_min = gr.Number(value=0, label="θ min")
                    theta_max = gr.Number(value=6.28, label="θ max")
            
            with gr.Group():
                gr.Markdown("### Style Settings")
                color = gr.ColorPicker(value="#1f77b4", label="Line Color")
                line_style = gr.Dropdown(
                    choices=["-", "--", "-.", ":"],
                    value="-",
                    label="Line Style"
                )
                color_scheme = gr.Dropdown(
                    choices=["viridis", "plasma", "inferno", "magma", "coolwarm", "RdYlBu"],
                    value="viridis",
                    label="3D Color Scheme"
                )
                show_grid = gr.Checkbox(value=True, label="Show Grid")
                custom_title = gr.Textbox(label="Custom Title (optional)")
            
            generate_btn = gr.Button("🎨 Generate Plot", variant="primary")
            
            # Example buttons
            gr.Markdown("### 📚 Examples")
            example_dropdown = gr.Dropdown(
                choices=["Select an example..."],
                label="Load Example"
            )
            
        with gr.Column(scale=2):
            output_plot = gr.Plot(label="Generated Plot")

    # Event handlers
    def update_inputs(plot_type):
        parametric_visible = plot_type == "Parametric (x = f(t), y = g(t))"
        
        # Update example dropdown
        if plot_type in examples:
            example_choices = ["Select an example..."] + [
                f"Example {i+1}" for i in range(len(examples[plot_type]))
            ]
        else:
            example_choices = ["Select an example..."]
        
        return (
            gr.update(visible=parametric_visible),
            gr.update(choices=example_choices, value="Select an example...")
        )

    def load_example_equations(plot_type, example_choice):
        if example_choice == "Select an example...":
            return "", "", "", ""
        
        try:
            example_idx = int(example_choice.split()[-1]) - 1
            return load_example(plot_type, example_idx)
        except:
            return "", "", "", ""

    plot_type.change(
        update_inputs,
        inputs=[plot_type],
        outputs=[parametric_inputs, example_dropdown]
    )

    example_dropdown.change(
        load_example_equations,
        inputs=[plot_type, example_dropdown],
        outputs=[equation, x_equation, y_equation, custom_title]
    )

    generate_btn.click(
        generate_plot,
        inputs=[
            plot_type, equation, x_equation, y_equation,
            x_min, x_max, y_min, y_max, t_min, t_max, theta_min, theta_max,
            color, line_style, color_scheme, show_grid, custom_title
        ],
        outputs=[output_plot]
    )

    # Auto-generate on equation change
    equation.change(
        generate_plot,
        inputs=[
            plot_type, equation, x_equation, y_equation,
            x_min, x_max, y_min, y_max, t_min, t_max, theta_min, theta_max,
            color, line_style, color_scheme, show_grid, custom_title
        ],
        outputs=[output_plot]
    )

if __name__ == "__main__":
    demo.launch(share=True)