import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
import matplotlib.animation
from sympy import latex, Matrix

def animate(exprs, var, 
            # --- Labels & Text ---
            labels=None, title="Animation", xlabel="x", ylabel="y",
            # --- Styling ---
            line_styles=None, colors=None, linewidth=2, 
            # --- Limits ---
            xlim=None, ylim=None, 
            # --- Animation Config ---
            animation_time=5000, frames=100, resolution=400, show=True):

    """
    Animates 2D curves from symbolic expressions or numeric datasets using Matplotlib.

    Parameters:
        exprs          : Expression, tuple, or list of expressions (symbolic, parametric, or numeric data)
        var            : Symbol or tuple defining variable/range, or array mode ('array', x_values)
        
        # --- Labels ---
        labels         : Curve labels (strings or SymPy expressions, LaTeX supported) (default=None)
        title          : Plot title (default="Animation")
        xlabel         : X-axis label (default="x")
        ylabel         : Y-axis label (default="y")
        
        # --- Styling ---
        line_styles    : Line styles for each curve (default=None -> solid)
        colors         : Colors for each curve (default=None -> Matplotlib cycle)
        linewidth      : Line width(s) as int or list (default=2)
        
        # --- Limits ---
        xlim           : X-axis limits as (min, max) (default=None -> auto)
        ylim           : Y-axis limits as (min, max) (default=None -> auto)
        
        # --- Animation Config ---
        animation_time : Total animation duration in milliseconds (default=5000)
        frames         : Number of animation frames (default=100)
        resolution     : Number of evaluation points for symbolic expressions (default=400)
        show           : Whether to display the animation (default=True)

    Returns:
        matplotlib.animation.FuncAnimation
            The animation object
    """

    # Configuration
    plt.rcParams["animation.html"] = "jshtml"
    plt.rcParams['figure.dpi'] = 150
    plt.ioff()

    # Label Formatting
    def _smart_label(lbl):
        """Converts strings or SymPy objects to LaTeX formatted strings."""
        if lbl is None:
            return None
        if isinstance(lbl, str):
            # Check for LaTeX characters
            if any(c in lbl for c in ['_', '^', '\\', '$']):
                return f"${lbl}$" if '$' not in lbl else lbl
            return lbl
        # Convert SymPy expression to LaTeX
        return f"${latex(lbl)}$"

    # Data Preparation
    def _prepare_data(raw_exprs, variable_def):
        """
        Parses inputs (Symbolic/Numeric) and generates consistent (x, y) arrays.
        """
        # Normalize inputs to lists
        if not isinstance(raw_exprs, list):
            expr_list = [raw_exprs]
        else:
            expr_list = raw_exprs

        # Determine symbol and range
        is_sympy_mode = True
        
        if isinstance(variable_def, tuple):
            if variable_def[0] == 'array':
                is_sympy_mode = False
                x_vals_sample = variable_def[1]
                # If array mode, frames usually equals array length
                nonlocal frames 
                frames = len(x_vals_sample) 
                x_sym = None
            else:
                x_sym = variable_def[0]
                x_range = variable_def[1]
                x_vals_sample = np.linspace(float(x_range[0]), float(x_range[1]), frames)
        else:
            x_sym = variable_def
            x_range = (-1, 1)
            x_vals_sample = np.linspace(float(x_range[0]), float(x_range[1]), frames)

        prepared_data = []

        for expr in expr_list:
            # Handle SymPy Matrices
            if isinstance(expr, Matrix):
                expr = np.array(expr).astype(np.float64).flatten()
            
            # Handle Parametric or Data Tuples (x, y)
            if isinstance(expr, (tuple, list)) and len(expr) == 2:
                x_comp, y_comp = expr
                
                # Symbolic Parametric
                if isinstance(x_comp, sp.Basic) or isinstance(y_comp, sp.Basic):
                    if is_sympy_mode:
                        f_x = sp.lambdify(x_sym, x_comp, modules='numpy')
                        f_y = sp.lambdify(x_sym, y_comp, modules='numpy')
                        x_data = f_x(x_vals_sample)
                        y_data = f_y(x_vals_sample)
                    else:
                         # Fallback if mixed types (shouldn't happen often)
                        x_data, y_data = x_comp, y_comp
                else:
                    # Numeric Data or Constants
                    x_data, y_data = x_comp, y_comp
                    
                    # specific check for constant arrays
                    is_x_array = hasattr(x_data, '__len__')
                    is_y_array = hasattr(y_data, '__len__')

                    if is_x_array and not is_y_array:
                        y_data = np.full_like(x_data, y_data)
                    elif not is_x_array and is_y_array:
                        x_data = np.full_like(y_data, x_data)

                prepared_data.append((np.array(x_data), np.array(y_data)))

            # Handle Single Expressions (y = f(x))
            else:
                expr = sp.sympify(expr)
                
                # Constant function check
                if is_sympy_mode and not expr.has(x_sym):
                    y_vals = np.full_like(x_vals_sample, float(expr))
                else:
                    if is_sympy_mode:
                        f = sp.lambdify(x_sym, expr, modules='numpy')
                        try:
                            y_vals = f(x_vals_sample)
                            # Handle case where lambdify returns a scalar for vector input
                            if np.isscalar(y_vals): 
                                y_vals = np.full_like(x_vals_sample, y_vals)
                        except Exception:
                            # Fallback for complex constant expressions
                            y_vals = np.full_like(x_vals_sample, float(expr))
                    else:
                        y_vals = np.array(expr).flatten()
                
                y_vals = np.array(y_vals).flatten()
                prepared_data.append((x_vals_sample, y_vals))
                
        return prepared_data

    # Animation Logic
    def _create_animation_object(plot_data):
        """Sets up the figure, axes, and FuncAnimation object."""
        
        # 1. Normalize Styling
        n_plots = len(plot_data)
        
        # Thickness
        if isinstance(linewidth, (int, float)):
            lw_list = [linewidth] * n_plots
        else:
            lw_list = linewidth
            
        # Limits Calculation (if None)
        nonlocal xlim, ylim
        if xlim is None:
            all_x = np.concatenate([d[0] for d in plot_data])
            x_min, x_max = np.min(all_x), np.max(all_x)
            dx = x_max - x_min if x_max != x_min else 1.0
            xlim = (x_min - 0.05 * dx, x_max + 0.05 * dx)
            
        if ylim is None:
            all_y = np.concatenate([d[1] for d in plot_data])
            y_min, y_max = np.min(all_y), np.max(all_y)
            dy = y_max - y_min if y_max != y_min else 1.0
            ylim = (y_min - 0.1 * dy, y_max + 0.1 * dy)

        # 2. Setup Figure
        fig, ax = plt.subplots(figsize=(10, 6))

        # 3. Update Function
        def _update_frame(frame):
            ax.clear()
            
            # Use 'frame + 1' to ensure we plot at least one point
            idx = int(frame + 1)
            
            for i, (x_data, y_data) in enumerate(plot_data):
                # Safe Styling Retrieval
                style = line_styles[i] if line_styles and i < len(line_styles) else 'solid'
                color = colors[i] if colors and i < len(colors) else None
                
                # Label Handling
                raw_lbl = labels[i] if labels and i < len(labels) else None
                lbl = _smart_label(raw_lbl)

                # Slice data for animation effect
                # Ensure we don't index out of bounds if frames > data length
                curr_x = x_data[:min(idx, len(x_data))]
                curr_y = y_data[:min(idx, len(y_data))]

                ax.plot(curr_x, curr_y,
                       label=lbl,
                       linestyle=style,
                       color=color,
                       linewidth=lw_list[i] if i < len(lw_list) else 2)

            # Apply Layout Settings
            ax.set_xlim(xlim)
            ax.set_ylim(ylim)
            ax.set_title(_smart_label(title))
            ax.set_xlabel(_smart_label(xlabel))
            ax.set_ylabel(_smart_label(ylabel))
            
            if labels:
                ax.legend()
            ax.grid(True, alpha=0.3)

        # 4. Create Animation
        interval = animation_time / frames
        anim = matplotlib.animation.FuncAnimation(
            fig, _update_frame, frames=frames, interval=interval, repeat=True
        )
        
        return anim

    # Main Execution
    
    # 1. Prepare Data
    data = _prepare_data(exprs, var)
    
    # 2. Create Animation
    anim_obj = _create_animation_object(data)
    
    # 3. Display
    if not show:
        plt.close()
        
    return anim_obj