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

def animate_arrow(position, vector, var,
                  # --- Background Path ---
                  path=None, showpath=False,
                  # --- Labels & Text ---
                  labels=None, title=None, xlabel=None, ylabel=None,
                  # --- Styling ---
                  colors=None, line_styles=None, linewidth=2,
                  scaling=None, fontsize=None,
                  # --- Arrow Shape ---
                  mutation_scale=25,
                  # --- Axis Options ---
                  center_axes=True, grid=True,
                  # --- Limits ---
                  xlim=None, ylim=None,
                  # --- Animation Config ---
                  animation_time=5000, frames=50, show=True,
                  # --- Figure Size ---
                  width=800, height=600):

    """
    Animates a moving arrow along a parametric path — mirrors Maple's:
        animate(arrow, [[C(1),C(2)], [dC], shape=arrow, color=green], frames=N, t=0..T)

    The arrow origin moves along position(t) while the arrow direction
    follows vector(t), both evaluated at each frame's parameter value.

    Parameters:
        position       : SymPy Matrix [x(t), y(t)] — arrow origin at each t
        vector         : SymPy Matrix [vx(t), vy(t)] — arrow direction at each t
        var            : (symbol, t_min, t_max) — parameter and range
                         e.g. (t, 0, 2)

        # --- Background Path ---
        path           : Optional SymPy Matrix [x(t), y(t)] to draw as static background
                         (default=None → uses position curve)
        showpath       : Whether to draw the background path curve (default=False)

        # --- Labels & Text ---
        labels         : List of legend labels — [arrow_label, path_label] (default=None)
        title          : Plot title (default=None → auto "f(t), t=t_min..t_max")
        xlabel         : X-axis label (default=None)
        ylabel         : Y-axis label (default=None)

        # --- Styling ---
        colors         : List of colors — [arrow_color, path_color] (default=None → green, blue)
        line_styles    : List of line styles — [path_style] (default=None → solid)
        linewidth      : Line width for arrow shaft and path (default=2)
        scaling        : 'constrained' → equal axis units (default=None)
        fontsize       : Base font size (default=None → auto)

        # --- Arrow Shape ---
        mutation_scale : Arrow head size in display points (default=25)
                         15=small, 25=medium, 35=large

        # --- Axis Options ---
        center_axes    : Move axes to origin (default=True)
        grid           : Show grid (default=True)

        # --- 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 duration in milliseconds (default=5000)
        frames         : Number of animation frames (default=50)
        show           : Display animation inline in Jupyter (default=True)

        # --- Figure Size ---
        width, height  : Figure size in pixels (default=800x600)

    Returns:
        matplotlib.animation.FuncAnimation
    """

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

    # Font Scaling
    def _get_font_sizes():
        if fontsize is None:
            scale    = min(width / 800, height / 600)
            title_fs = max(4, 16 * scale)
            label_fs = max(4, 14 * scale)
            tick_fs  = max(4, 12 * scale)
            leg_fs   = max(4, 12 * scale)
        else:
            title_fs = fontsize + 2
            label_fs = fontsize
            tick_fs  = max(4, fontsize - 2)
            leg_fs   = fontsize
        return title_fs, label_fs, tick_fs, leg_fs

    # Lambdify a SymPy Matrix component-wise → list of callables
    def _lambdify_vec(mat, t_sym):
        return [sp.lambdify(t_sym, mat[i], modules='numpy') for i in range(len(mat))]

    # Evaluate lambdified list at array → (2, N) numpy array
    def _eval_vec(funcs, t_arr):
        rows = []
        for f in funcs:
            v = f(t_arr)
            if np.isscalar(v):
                v = np.full_like(t_arr, float(v))
            rows.append(np.asarray(v, dtype=float))
        return np.array(rows)   # shape (2, N)

    # Compute auto xlim / ylim from all data
    def _compute_limits(pos_f, vec_f, path_f, t_arr):
        pos  = _eval_vec(pos_f, t_arr)
        vec  = _eval_vec(vec_f, t_arr)

        all_x = np.concatenate([pos[0], pos[0] + vec[0]])
        all_y = np.concatenate([pos[1], pos[1] + vec[1]])

        if showpath and path_f is not None:
            p = _eval_vec(path_f, t_arr)
            all_x = np.concatenate([all_x, p[0]])
            all_y = np.concatenate([all_y, p[1]])

        fx = all_x[np.isfinite(all_x)]
        fy = all_y[np.isfinite(all_y)]

        pad_x = (fx.max() - fx.min()) * 0.2 if fx.max() != fx.min() else 0.5
        pad_y = (fy.max() - fy.min()) * 0.2 if fy.max() != fy.min() else 0.5

        return (fx.min() - pad_x, fx.max() + pad_x), \
               (fy.min() - pad_y, fy.max() + pad_y)

    # Setup Figure & Axes
    def _setup_figure(xlim_, ylim_, title_str):
        title_fs, label_fs, tick_fs, _ = _get_font_sizes()

        fig, ax = plt.subplots(figsize=(width / 100, height / 100))
        ax.set_xlim(xlim_)
        ax.set_ylim(ylim_)
        ax.grid(grid, alpha=0.3)
        ax.set_title(_smart_label(title_str), fontsize=title_fs)
        ax.tick_params(axis='both', labelsize=tick_fs)

        if scaling:
            ax.set_aspect('equal', adjustable='box')

        if center_axes:
            ax.spines['left'].set_position('zero')
            ax.spines['bottom'].set_position('zero')
            ax.spines['right'].set_visible(False)
            ax.spines['top'].set_visible(False)
            ax.xaxis.set_ticks_position('bottom')
            ax.yaxis.set_ticks_position('left')
            ax.tick_params(axis='x', pad=5)
            ax.tick_params(axis='y', pad=5)
            if xlabel:
                ax.annotate(_smart_label(xlabel),
                            xy=(1, 0), xycoords=('axes fraction', 'data'),
                            xytext=(10, 0), textcoords='offset points',
                            ha='left', va='center', fontsize=label_fs)
            if ylabel:
                ax.annotate(_smart_label(ylabel),
                            xy=(0, 1), xycoords=('data', 'axes fraction'),
                            xytext=(0, 10), textcoords='offset points',
                            ha='center', va='bottom', fontsize=label_fs)
        else:
            if xlabel: ax.set_xlabel(_smart_label(xlabel), fontsize=label_fs)
            if ylabel: ax.set_ylabel(_smart_label(ylabel), fontsize=label_fs)

        return fig, ax

    # MAIN EXECUTION

    plt.rcParams["animation.html"]        = "jshtml"
    plt.rcParams['figure.dpi']            = 100
    plt.rcParams['animation.embed_limit'] = 50.0
    plt.ioff()

    # Unpack variable
    def _unpack_var(var):
        if len(var) == 2:
            t_sym, t_range = var
            return t_sym, float(t_range[0]), float(t_range[1])
        elif len(var) == 3:
            t_sym, t_min, t_max = var
            return t_sym, float(t_min), float(t_max)
        else:
            raise ValueError("var must be (symbol, (t_min, t_max)) or (symbol, t_min, t_max)")

    t_sym, t_min_f, t_max_f = _unpack_var(var)
    t_vals  = np.linspace(t_min_f, t_max_f, frames)

    # Colors
    arrow_color = colors[0] if colors and len(colors) > 0 else 'green'
    path_color  = colors[1] if colors and len(colors) > 1 else 'blue'

    # Lambdify
    pos_f  = _lambdify_vec(position, t_sym)
    vec_f  = _lambdify_vec(vector,   t_sym)
    path_f = _lambdify_vec(path if path is not None else position, t_sym)

    # Auto title
    auto_title = title or f"$t = {t_min_f:.4f}$"

    # Compute limits
    _xlim, _ylim = _compute_limits(pos_f, vec_f, path_f, t_vals)
    _xlim = xlim or _xlim
    _ylim = ylim or _ylim

    # Setup figure
    fig, ax = _setup_figure(_xlim, _ylim, auto_title)
    _, _, _, leg_fs = _get_font_sizes()

    # Title text object for live t update
    title_fs, _, _, _ = _get_font_sizes()
    title_text = ax.set_title(f"$t = {t_min_f:.4f}$", fontsize=title_fs)

    # Draw static background path
    if showpath:
        p      = _eval_vec(path_f, t_vals)
        p_lbl  = _smart_label(labels[1]) if labels and len(labels) > 1 else None
        p_sty  = line_styles[0] if line_styles and len(line_styles) > 0 else '-'
        ax.plot(p[0], p[1],
                color     = path_color,
                linestyle = p_sty,
                linewidth = linewidth,
                label     = p_lbl,
                zorder    = 1)

    # Arrow legend proxy
    arrow_lbl = _smart_label(labels[0]) if labels and len(labels) > 0 else None
    if arrow_lbl:
        ax.plot([], [], color=arrow_color, linewidth=linewidth, label=arrow_lbl)

    if labels:
        ax.legend(fontsize=leg_fs)

    # Arrow artists — shaft line + head annotation updated each frame
    shaft_line, = ax.plot([], [], color=arrow_color, lw=linewidth,
                          solid_capstyle='butt', zorder=2)
    arrow_head  = [None]   # list allows reassignment inside nested function

    # Update Function
    def _update_frame(frame):
        t_val = t_vals[frame]

        ox = float(pos_f[0](t_val))
        oy = float(pos_f[1](t_val))
        vx = float(vec_f[0](t_val))
        vy = float(vec_f[1](t_val))

        # Update shaft
        shaft_line.set_data([ox, ox + vx], [oy, oy + vy])

        # Remove old head, draw new one at tip
        if arrow_head[0] is not None:
            arrow_head[0].remove()

        arrow_head[0] = ax.annotate('',
            xy     = (ox + vx,               oy + vy              ),
            xytext = (ox + vx - vx * 0.001,  oy + vy - vy * 0.001),
            arrowprops = dict(
                arrowstyle     = '-|>',
                mutation_scale = mutation_scale,
                color          = arrow_color,
                lw             = linewidth,
                shrinkA        = 0,
                shrinkB        = 0,
            ),
            zorder = 3
        )

        # Live t value in title  (matches Maple "t = 2.0000")
        title_text.set_text(f"$t = {t_val:.4f}$")

        return [shaft_line]

    # Create animation
    interval = animation_time / max(frames, 1)
    anim = FuncAnimation(fig, _update_frame, frames=frames,
                         interval=interval, repeat=True, blit=False)

    if show:
        ipy_display(HTML(anim.to_jshtml()))

    return anim