import numpy as np
import plotly.graph_objects as go
from sympy import symbols, lambdify
from sympy.parsing.sympy_parser import parse_expr, standard_transformations, implicit_multiplication_application
from solver import preprocess_equation

def plot_function(input_str):
    """Create an interactive plot of the function."""
    try:
        # Preprocess input
        processed_input = preprocess_equation(input_str)

        # Handle different types of input
        if '=' in processed_input:
            # Equation: move everything to left side
            left_side, right_side = [side.strip() for side in processed_input.split('=')]
            transformations = standard_transformations + (implicit_multiplication_application,)
            expr = parse_expr(left_side, transformations=transformations) - parse_expr(right_side, transformations=transformations)
        elif input_str.startswith('∫'):
            # Integration: plot the original function
            expr = parse_expr(processed_input[1:].strip(), transformations=(standard_transformations + (implicit_multiplication_application,)))
        elif input_str.startswith('d/dx'):
            # Derivative: plot the original function
            expr = parse_expr(processed_input[4:].strip(), transformations=(standard_transformations + (implicit_multiplication_application,)))
        else:
            # Regular expression
            expr = parse_expr(processed_input, transformations=(standard_transformations + (implicit_multiplication_application,)))

        # Create lambda function for numpy evaluation
        x = symbols('x')
        f = lambdify(x, expr, 'numpy')

        # Generate x values
        x_vals = np.linspace(-10, 10, 1000)

        # Calculate y values
        y_vals = f(x_vals)

        # Create plot
        fig = go.Figure()

        # Add function curve
        fig.add_trace(go.Scatter(
            x=x_vals,
            y=y_vals,
            mode='lines',
            name='f(x)',
            line=dict(color='#FF4B4B', width=2)
        ))

        # Add x-axis line
        fig.add_trace(go.Scatter(
            x=x_vals,
            y=[0]*len(x_vals),
            mode='lines',
            name='x-axis',
            line=dict(color='black', width=1)
        ))

        # Update layout
        fig.update_layout(
            title='Function Visualization',
            xaxis_title='x',
            yaxis_title='y',
            showlegend=True,
            hovermode='x',
            plot_bgcolor='white',
            width=800,
            height=500
        )

        # Update axes
        fig.update_xaxes(zeroline=True, zerolinewidth=1, zerolinecolor='black', gridcolor='lightgray')
        fig.update_yaxes(zeroline=True, zerolinewidth=1, zerolinecolor='black', gridcolor='lightgray')

        return fig

    except Exception as e:
        raise Exception(f"Error creating plot: {str(e)}")