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, HTML

def animate(exprs, var, **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
    """

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

    # Unpack Options
    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)
    scaling = kwargs.get('scaling', None)
    
    # Styles
    labels = kwargs.get('labels', [])
    colors = kwargs.get('colors', ['black', 'blue', 'red', 'green'])
    line_styles = kwargs.get('line_styles', ['solid'] * 10)
    linewidth = kwargs.get('linewidth', 2) 
    
    # Animation Settings
    animation_time = kwargs.get('animation_time', 5000)
    frames = kwargs.get('frames', 100)
    show = kwargs.get('show', True)
    width = kwargs.get('width', 800)
    height = kwargs.get('height', 600)

    # Helper: Range Unpacker
    def _unpack_range(r_tuple):
        if len(r_tuple) == 3: return r_tuple[0], float(r_tuple[1]), float(r_tuple[2])
        elif len(r_tuple) == 2: return r_tuple[0], float(r_tuple[1][0]), float(r_tuple[1][1])
        raise ValueError(f"Invalid range format: {r_tuple}")

    # Parse Main Animation Variable
    t_sym, t_start, t_end = _unpack_range(var)
    t_vals = np.linspace(t_start, t_end, frames)

    # Process Expressions & Generate Frame Data
    all_curves_data = [] 

    if not isinstance(exprs, list): exprs = [exprs]

    for i, expr in enumerate(exprs):
        
        # Check for Deformation Mode: (x, y, (z, 0, h))
        if isinstance(expr, (tuple, list)) and len(expr) == 3 and isinstance(expr[2], (tuple, list)):
            x_e, y_e, spatial_info = expr
            z_sym, z_min, z_max = _unpack_range(spatial_info)
            
            # Create Spatial Mesh
            z_vals = np.linspace(z_min, z_max, 200) 
            
            # Lambdify f(z, t)
            fx = sp.lambdify((z_sym, t_sym), x_e, modules='numpy')
            fy = sp.lambdify((z_sym, t_sym), y_e, modules='numpy')
            
            frame_x_list = []
            frame_y_list = []
            
            for t in t_vals:
                x_res = fx(z_vals, t)
                y_res = fy(z_vals, t)
                
                if not hasattr(x_res, '__len__'): x_res = np.full_like(z_vals, x_res)
                if not hasattr(y_res, '__len__'): y_res = np.full_like(z_vals, y_res)
                
                frame_x_list.append(x_res)
                frame_y_list.append(y_res)
                
            all_curves_data.append({
                'type': 'deformation',
                'x': frame_x_list, 
                'y': frame_y_list
            })

        # Standard Trace Mode
        else:
            if isinstance(expr, (tuple, list)) and len(expr) == 2:
                x_e, y_e = expr
            else:
                x_e, y_e = t_sym, expr
            
            fx = sp.lambdify(t_sym, x_e, modules='numpy')
            fy = sp.lambdify(t_sym, y_e, modules='numpy')
            
            full_x = fx(t_vals)
            full_y = fy(t_vals)
            
            if not hasattr(full_x, '__len__'): full_x = np.full_like(t_vals, full_x)
            if not hasattr(full_y, '__len__'): full_y = np.full_like(t_vals, full_y)
            
            all_curves_data.append({
                'type': 'trace',
                'x': full_x, 
                'y': full_y
            })

    # Auto-Scaling
    if xlim is None or ylim is None:
        flat_x, flat_y = [], []
        for c in all_curves_data:
            if c['type'] == 'deformation':
                flat_x.append(np.concatenate(c['x']))
                flat_y.append(np.concatenate(c['y']))
            else:
                flat_x.append(c['x'])
                flat_y.append(c['y'])
        
        flat_x = np.concatenate(flat_x)
        flat_y = np.concatenate(flat_y)
        
        if xlim is None:
            xmin, xmax = np.min(flat_x), np.max(flat_x)
            pad = (xmax - xmin) * 0.1 if xmax != xmin else 1.0
            xlim = (xmin - pad, xmax + pad)
            
        if ylim is None:
            ymin, ymax = np.min(flat_y), np.max(flat_y)
            pad = (ymax - ymin) * 0.1 if ymax != ymin else 1.0
            ylim = (ymin - pad, ymax + pad)

    # Plotting Setup
    fig, ax = plt.subplots(figsize=(width/100, height/100))
    
    if scaling == 'constrained' or (isinstance(scaling, list) and 'constrained' in scaling):
        ax.set_aspect('equal')
        
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.grid(True, alpha=0.3)
    
    def fmt_latex(l): 
        return f"${latex(l)}$" if hasattr(l, 'free_symbols') or (isinstance(l, str) and '\\' in l) else str(l)
    
    ax.set_title(fmt_latex(title))
    ax.set_xlabel(fmt_latex(xlabel))
    ax.set_ylabel(fmt_latex(ylabel))

    # Initialize Lines
    lines = []
    for i in range(len(all_curves_data)):
        col = colors[i % len(colors)]
        sty = line_styles[i % len(line_styles)]
        lbl = labels[i] if i < len(labels) else None
        
        if isinstance(linewidth, list): lw = linewidth[i % len(linewidth)]
        else: lw = linewidth
            
        ln, = ax.plot([], [], color=col, linestyle=sty, linewidth=lw, label=fmt_latex(lbl) if lbl else None)
        lines.append(ln)

    if any(labels): ax.legend()

    # Animation Loop
    def update_frame(frame_idx):
        t_curr = t_vals[frame_idx]
        ax.set_title(f"{fmt_latex(title)} ($t={t_curr:.2f}$)")
        
        for i, data in enumerate(all_curves_data):
            if data['type'] == 'deformation':
                lines[i].set_data(data['x'][frame_idx], data['y'][frame_idx])
            else:
                lines[i].set_data(data['x'][:frame_idx+1], data['y'][:frame_idx+1])
                
    ani = matplotlib.animation.FuncAnimation(
        fig, update_frame, frames=frames, interval=animation_time/frames, blit=False
    )
    
    if show:
        display(HTML(ani.to_jshtml()))
        plt.close()
        
    return ani