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

def animate(exprs, var=None, **kwargs):
    
    """
    Animates 2D curves from symbolic expressions or numeric datasets using Matplotlib.

    Parameters:
        exprs          : Expression, tuple, or list of expressions to animate (symbolic, parametric, or numeric data)
        var            : Symbol or tuple defining the variable and range, or array mode ('array', x_values)
        labels         : Curve labels (strings or SymPy expressions, LaTeX supported) (default=None)
        line_styles    : Line styles for each curve (default=None → solid)
        colors         : Colors for each curve (default=None → Matplotlib cycle)
        title          : Plot title (string or SymPy expression, LaTeX supported) (default="Animation")
        xlabel         : X-axis label (string or SymPy expression, LaTeX supported) (default="x")
        ylabel         : Y-axis label (string or SymPy expression, LaTeX supported) (default="y")
        xlim           : X-axis limits as (min, max) (default=None → auto)
        ylim           : Y-axis limits as (min, max) (default=None → auto)
        linewidth      : Line width(s) as int or list (default=2)
        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)
        width          : Plot width in pixels (default=800)
        height         : Plot height in pixels (default=600)

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

    # Unpack Options (The "Unfolding")
    labels = kwargs.get('labels', None)
    line_styles = kwargs.get('line_styles', None)
    colors = kwargs.get('colors', None)
    
    title = kwargs.get('title', "Animation")
    xlabel = kwargs.get('xlabel', "x")
    ylabel = kwargs.get('ylabel', "y")
    
    xlim = kwargs.get('xlim', None)
    ylim = kwargs.get('ylim', None)
    
    thickness = kwargs.get('thickness', 2)
    scaling = kwargs.get('scaling', None)
    
    animation_time = kwargs.get('animation_time', 5000)
    frames = kwargs.get('frames', 100)
    resolution = kwargs.get('resolution', 400)
    show = kwargs.get('show', True)
    
    width = kwargs.get('width', 800)
    height = kwargs.get('height', 600)

    # Smart Labels
    def _smart_label(lbl):
        if isinstance(lbl, str):
            if any(c in lbl for c in ['_', '^', '\\']): return f"${lbl}$"
            return lbl
        elif lbl is None: return None
        return f"${latex(lbl)}$"

    # Extract Data from Existing Plot
    def _extract_from_axes(ax_input):
        extracted_data, extracted_styles, extracted_colors = [], [], []
        extracted_labels, extracted_widths = [], []
        for line in ax_input.lines:
            extracted_data.append((line.get_xdata(), line.get_ydata()))
            extracted_colors.append(line.get_color())
            extracted_styles.append(line.get_linestyle())
            extracted_labels.append(line.get_label())
            extracted_widths.append(line.get_linewidth())
        return extracted_data, extracted_styles, extracted_colors, extracted_labels, extracted_widths

    # Handle Parametric Logic
    def _get_parametric_data(expr, x_sym, frames_count):
        if len(expr) == 3: # Custom Range
            x_e, y_e, range_info = expr
            t_sym, t_min, t_max = range_info
            t_vals = np.linspace(float(t_min), float(t_max), frames_count)
            f_x = sp.lambdify(t_sym, x_e, modules='numpy')
            f_y = sp.lambdify(t_sym, y_e, modules='numpy')
            x_res, y_res = f_x(t_vals), f_y(t_vals)
            if not hasattr(x_res, '__len__'): x_res = np.full_like(t_vals, x_res)
            if not hasattr(y_res, '__len__'): y_res = np.full_like(t_vals, y_res)
            return np.array(x_res).flatten(), np.array(y_res).flatten()
        elif len(expr) == 2: # Global Range
            x_data, y_data = expr
            if isinstance(x_data, sp.Basic) or isinstance(y_data, sp.Basic):
                f_x = sp.lambdify(x_sym, x_data, modules='numpy')
                f_y = sp.lambdify(x_sym, y_data, modules='numpy')
                x_data, y_data = f_x(x_vals_sample), f_y(x_vals_sample)
                if hasattr(x_data, '__len__') and not hasattr(y_data, '__len__'):
                    y_data = np.full_like(x_data, y_data)
                elif not hasattr(x_data, '__len__') and hasattr(y_data, '__len__'):
                    x_data = np.full_like(y_data, x_data)
            return x_data, y_data
        return None, None

    # Handle Standard y=f(x) Logic
    def _get_standard_data(expr, x_sym, x_vals_sample):
        expr = sp.sympify(expr)
        if not expr.has(x_sym):
            y_vals = np.full_like(x_vals_sample, float(expr))
        else:
            f = sp.lambdify(x_sym, expr, modules='numpy')
            y_vals = np.array(f(x_vals_sample)).flatten()
        return x_vals_sample, y_vals

    # Main Execution
    dpi = 100 
    plt.rcParams["animation.html"] = "jshtml"
    plt.rcParams['figure.dpi'] = dpi
    plt.ioff()

    plot_data = []
    final_colors, final_styles, final_labels, final_widths = [], [], [], []

    # Existing Plot
    if isinstance(exprs, plt.Axes):
        plot_data, final_styles, final_colors, final_labels, final_widths = _extract_from_axes(exprs)
        if plot_data: frames = max(len(d[0]) for d in plot_data)
        if title == "Animation": title = exprs.get_title()
        if xlabel == "x": xlabel = exprs.get_xlabel()
        if ylabel == "y": ylabel = exprs.get_ylabel()

    # List of Expressions
    else:
        if not isinstance(exprs, list): exprs = [exprs]
        
        # Parse Variable/Range
        if isinstance(var, tuple):
            if var[0] == 'array':
                x_sym, x_vals_sample = None, var[1]
                frames = len(x_vals_sample)
            else:
                x_sym, x_range = var[0], var[1]
                x_vals_sample = np.linspace(float(x_range[0]), float(x_range[1]), frames)
        else:
            x_sym = var
            x_range = (-1, 1)
            x_vals_sample = np.linspace(float(x_range[0]), float(x_range[1]), frames)

        # Process Expressions
        for i, expr in enumerate(exprs):
            if isinstance(expr, Matrix): expr = np.array(expr).astype(np.float64).flatten()
            
            # Style defaults
            usr_style = line_styles[i] if line_styles and i < len(line_styles) else None
            usr_color = colors[i] if colors and i < len(colors) else None
            usr_label = labels[i] if labels and i < len(labels) else None
            usr_width = thickness[i] if isinstance(thickness, list) and i < len(thickness) else (thickness if isinstance(thickness, (int, float)) else 2)

            final_styles.append(usr_style)
            final_colors.append(usr_color)
            final_labels.append(usr_label)
            final_widths.append(usr_width)

            if isinstance(expr, (tuple, list)):
                px, py = _get_parametric_data(expr, x_sym, frames)
                plot_data.append((px, py))
            else:
                sx, sy = _get_standard_data(expr, x_sym, x_vals_sample)
                plot_data.append((sx, sy))

    # Limits & Figure Setup
    if xlim is None and plot_data:
        all_x = np.concatenate([d[0] for d in plot_data])
        min_x, max_x = np.min(all_x), np.max(all_x)
        pad_x = (max_x - min_x) * 0.1 if max_x != min_x else 0.5
        xlim = (min_x - pad_x, max_x + pad_x)

    if ylim is None and plot_data:
        all_y = np.concatenate([d[1] for d in plot_data])
        min_y, max_y = np.min(all_y), np.max(all_y)
        pad_y = (max_y - min_y) * 0.1 if max_y != min_y else 0.5
        ylim = (min_y - pad_y, max_y + pad_y)

    fig, ax = plt.subplots(figsize=(width / dpi, height / dpi))
    if scaling == 'constrained': ax.set_aspect('equal', adjustable='box')

    def animate_frame(frame):
        ax.clear()
        for i, (x_d, y_d) in enumerate(plot_data):
            total_points = len(x_d)
            idx = int((frame + 1) / frames * total_points)
            lbl = _smart_label(final_labels[i])
            stl = final_styles[i] if final_styles[i] else 'solid'
            clr = final_colors[i]
            wdth = final_widths[i]
            ax.plot(x_d[:idx], y_d[:idx], label=lbl, linestyle=stl, color=clr, linewidth=wdth)
        
        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 any(final_labels): ax.legend(loc='upper right')
        ax.grid(True, alpha=0.3)

    anim = matplotlib.animation.FuncAnimation(
        fig, animate_frame, frames=frames, interval=animation_time / frames, repeat=True
    )
    
    if show:
        display(anim)
        plt.close()
    
    return anim