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=1, scaling=None, fontsize=None,
            # --- Limits ---
            xlim=None, ylim=None, 
            # --- Animation Config ---
            animation_time=5000, frames=100, show=True,
            # --- Figure Size ---
            width=800, height=600):

    """
    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'] = 300
    plt.rcParams['animation.embed_limit'] = 50.0
    plt.ioff()

    # Label Formatting
    def _smart_label(lbl):
<<<<<<< HEAD
        """Converts strings or SymPy objects to LaTeX formatted strings."""
        if lbl is None:
            return None
            return lbl
>>>>>>> dev
        return f"${latex(lbl)}$"

    # Data Preparation
<<<<<<< HEAD
        """
        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
        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:
            if isinstance(expr, Matrix):
                expr = np.array(expr).astype(np.float64).flatten()
            
            if isinstance(expr, (tuple, list)) and len(expr) == 2:
                x_comp, y_comp = expr
                
                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:
                        x_data, y_data = x_comp, y_comp
                else:
                    x_data, y_data = x_comp, y_comp
                    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)))

            else:
                expr = sp.sympify(expr)
>>>>>>> dev
                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)
                            if np.isscalar(y_vals): y_vals = np.full_like(x_vals_sample, y_vals)
                            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):
        n_plots = len(plot_data)
        
        if isinstance(linewidth, (int, float)): lw_list = [linewidth] * n_plots
        else: lw_list = linewidth
            
        nonlocal xlim, ylim
        if xlim is None:
            all_x = np.concatenate([d[0] for d in plot_data])
            x_min, x_max = np.nanmin(all_x), np.nanmax(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.nanmin(all_y), np.nanmax(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=(width / 100, height / 100))
        
        # Dynamic Font Scaling
        if fontsize is None:
            # Scale relative to a standard 800x600 resolution, keeping a minimum floor 
            # so text never drops below 4pt and disappears completely.
            font_scale = min(width / 800, height / 600)
            title_fs = max(4, 16 * font_scale)
            label_fs = max(4, 14 * font_scale)
            tick_fs  = max(4, 12 * font_scale)
            legend_fs = max(4, 12 * font_scale)
        else:
            # If the user explicitly sets `fontsize`, use that as a baseline.
            title_fs = fontsize + 2
            label_fs = fontsize
            tick_fs  = max(4, fontsize - 2)
            legend_fs = fontsize
        
        # Apply Static Configuration ONCE
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        ax.set_title(_smart_label(title), fontsize=title_fs)
        ax.set_xlabel(_smart_label(xlabel), fontsize=label_fs)
        ax.set_ylabel(_smart_label(ylabel), fontsize=label_fs)
        ax.tick_params(axis='both', which='major', labelsize=tick_fs)
        ax.grid(True, alpha=0.3)
        
        if scaling:
            if isinstance(scaling, str):
                scaling_key = scaling.lower()
            else:
                scaling_key = 'equal'

            if scaling_key in ('equal', 'constrained', 'box', 'datalim'):
                x_span = xlim[1] - xlim[0]
                y_span = ylim[1] - ylim[0]
                if x_span > 0 and y_span > 0 and x_span != y_span:
                    if x_span < y_span:
                        pad = (y_span - x_span) / 2.0
                        xlim = (xlim[0] - pad, xlim[1] + pad)
                        ax.set_xlim(xlim)
                    else:
                        pad = (x_span - y_span) / 2.0
                        ylim = (ylim[0] - pad, ylim[1] + pad)
                        ax.set_ylim(ylim)

                adjustable = 'box' if scaling_key in ('equal', 'constrained') else scaling_key
                ax.set_aspect('equal', adjustable=adjustable)

        lines = []
        for i in range(n_plots):
            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
            raw_lbl = labels[i] if labels and i < len(labels) else None
            lbl = _smart_label(raw_lbl)
            lw = lw_list[i] if i < len(lw_list) else 2
            
            line, = ax.plot([], [], label=lbl, linestyle=style, color=color, linewidth=lw)
            lines.append(line)

        if labels: 
            # Apply dynamic font size to legend
            ax.legend(fontsize=legend_fs)

        # 3. Update Function
        def _update_frame(frame):
            idx = int(frame + 1)
            
            for i, (x_data, y_data) in enumerate(plot_data):
                curr_x = x_data[:min(idx, len(x_data))]
                curr_y = y_data[:min(idx, len(y_data))]
                lines[i].set_data(curr_x, curr_y)

            return lines

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

    # Main Execution
    data = _prepare_data(exprs, var)
    anim_obj = _create_animation_object(data)
    
    if not show:
        plt.close()
        
    return anim_obj