import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import plotly.graph_objects as go
from   matplotlib.patches import FancyArrow
from   sympy import latex

def plot_arrow(vectors, origins=None,
               # --- Plot Labels ---
               title=None, xlabel=None, ylabel=None, zlabel=None,
               labels=None,
               # --- Styling ---
               colors=None, linewidth=2, scaling='constrained', grid=True,
               # --- Arrow Shape ---
               shape='double_arrow', length=None, width=None, head_length=None, head_width=None, 
               # --- Backend ---
               backend=None,  
               # --- Typography ---
               fontsize=14, title_size=None, label_size=None, tick_size=None, legend_size=None,
               # --- Axis Options ---
               center_axes=True,
               # --- View angle (matplotlib only) ---
               elev=20, azim=45,
               # --- Limits ---
               xlim=None, ylim=None, zlim=None,
               # --- Display settings ---
               show=True, width_px=800, height_px=600):

    r"""
    Plots 2D or 3D vectors as arrows — mirrors Maple's plots[arrow]() command.

    Automatically detects 2D vs 3D from vector dimension.
    Backend is auto-selected: matplotlib for 2D, plotly for 3D (override with backend=).

    Parameters:
        vectors        : List of vectors — SymPy Matrix, list, or numpy array
                         2D: [x, y]      e.g. sp.Matrix([1, 2])
                         3D: [x, y, z]   e.g. sp.Matrix([1, 2, 3])
        origins        : List of origin points (default=None → all at zero)

        # --- Plot Labels ---
        labels         : Legend labels for each vector (default=None)
        title          : Plot title (default=None)
        xlabel         : X-axis label (default=None)
        ylabel         : Y-axis label (default=None)
        zlabel         : Z-axis label — 3D only (default=None)

        # --- Styling ---
        colors         : Colors for each arrow (default=None → cycle)
        linewidth      : Shaft line width for harpoon/arrow shapes (default=2)
        scaling        : 'constrained' → equal axis units — 2D matplotlib only (default='constrained')
        grid           : Whether to display a grid (default=True)

        # --- Arrow Shape  ---
        shape          : 'harpoon'          → line shaft + single angled line head (2D/3D)
                         'arrow'            → line shaft + two lines forming V head (2D/3D)
                         'double_arrow'     → filled rectangle shaft + filled triangle head (2D default)
                         'cylindrical_arrow'→ cylinder shaft + cone head (3D default, plotly only)

        # --- Arrow Dimensions  ---
        length         : Total arrow length override
                         - float: absolute length
                         - [float, relative=True]: length as fraction of vector magnitude
                         - None: use full vector magnitude (default)
        width          : Shaft width
                         - float: absolute width
                         - [float, relative=True]: width relative to arrow length
                         - None: 1/20 of arrow length (Maple default)
        head_length    : Length of arrowhead
                         - float: absolute head length
                         - [float, relative=True]: relative to arrow length
                         - None: 1/5 of arrow length (Maple default)
        head_width     : Width of arrowhead
                         - float: absolute head width
                         - [float, relative=True]: relative to shaft width
                         - None: 2x shaft width (Maple default)

        # --- Backend ---
        backend        : 'matplotlib' | 'plotly' | None=auto
                         Auto: matplotlib for 2D, plotly for 3D
                         'cylindrical_arrow' forces plotly for 3D

        # --- Typography ---
        fontsize       : Base font size (default=14)
        title_size     : Font size for title (default=None → fontsize+2)
        label_size     : Font size for axis labels (default=None → fontsize)
        tick_size      : Font size for tick labels (default=None → fontsize-2)
        legend_size    : Font size for legend (default=None → tick_size)

        # --- Axis Options ---
        center_axes    : matplotlib 2D only — move axes to origin (default=True)

        # --- View Angle (matplotlib 3D only) ---
        elev           : Elevation angle in degrees (default=20)
                         0° = horizontal view, 90° = top-down view
        azim           : Azimuth angle in degrees (default=45)
                         0° = front view, 90° = side view, 180° = back view

        # --- Limits ---
        xlim           : X-axis limits as (min, max) (default=None → auto)
        ylim           : Y-axis limits as (min, max) (default=None → auto)
        zlim           : Z-axis limits as (min, max) — 3D only (default=None → auto)

        # --- Display settings ---
        show           : Whether to display the plot (default=True)
        width_px       : Figure width in pixels (default=800)
        height_px      : Figure height in pixels (default=600)

    Returns:
        matplotlib Axes        (2D)
        matplotlib Axes3D      (3D, backend='matplotlib')
        plotly Figure          (3D, backend='plotly' or auto)
    """

    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)}$"

    def _to_vec(v):
        if isinstance(v, sp.MatrixBase):
            return np.array(v.tolist(), dtype=float).flatten()
        if isinstance(v, (list, tuple)):
            return np.array(v, dtype=float).flatten()
        if isinstance(v, np.ndarray):
            return v.astype(float).flatten()
        raise TypeError(f"Cannot convert {type(v)} to a vector.")

    def _auto_limits(tips, origins_np):
        all_pts = np.array(tips + origins_np)
        mins    = all_pts.min(axis=0)
        maxs    = all_pts.max(axis=0)
        spans   = maxs - mins
        pads    = np.where(spans > 0, spans * 0.2, 0.5)
        return list(zip(mins - pads, maxs + pads))

    def _get_font_sizes():
        _ts  = title_size  or fontsize + 2
        _ls  = label_size  or fontsize
        _tks = tick_size   or fontsize - 2
        _lgs = legend_size or _tks
        return _ts, _ls, _tks, _lgs

    def _resolve_dim(param, ref_length, default_abs):
        """
        Resolve a length/width parameter that can be:
          - None            → use default_abs
          - float           → absolute value
          - [float, True]   → relative to ref_length
          - [float, False]  → absolute value
        """
        if param is None:
            return default_abs
        if isinstance(param, (int, float)):
            return float(param)
        if isinstance(param, (list, tuple)) and len(param) == 2:
            val, relative = param
            return float(val) * ref_length if relative else float(val)
        return default_abs

    def _compute_arrow_dims(mag):
        """Compute all arrow dimensions from the Maple defaults and user overrides."""
        arr_len  = _resolve_dim(length,      mag,     mag)
        w        = _resolve_dim(width,       arr_len, arr_len / 20.0)
        hl       = _resolve_dim(head_length, arr_len, arr_len / 5.0)
        hw       = _resolve_dim(head_width,  w,       w * 2.0)
        return arr_len, w, hl, hw

    def _draw_harpoon_2d(ax, ox, oy, vx, vy, color, lw, lbl, arr_len, hl):
        """Line shaft + single angled line for head."""
        mag   = np.sqrt(vx**2 + vy**2)
        scale = arr_len / mag if mag > 0 else 1.0
        ex, ey = ox + vx * scale, oy + vy * scale

        ax.plot([ox, ex], [oy, ey], color=color, lw=lw, label=lbl,
                zorder=2, solid_capstyle='butt')

        if mag > 0:
            ux, uy = vx / mag, vy / mag
            px, py = -uy, ux
            hx = ex - ux * hl
            hy = ey - uy * hl
            hw_single = hl * 0.5
            ax.plot([hx + px * hw_single, ex],
                    [hy + py * hw_single, ey],
                    color=color, lw=lw, zorder=3)

    def _draw_arrow_2d_shape(ax, ox, oy, vx, vy, color, lw, lbl, arr_len, hl):
        """Line shaft + two lines forming a V head."""
        mag   = np.sqrt(vx**2 + vy**2)
        scale = arr_len / mag if mag > 0 else 1.0
        ex, ey = ox + vx * scale, oy + vy * scale

        ax.plot([ox, ex], [oy, ey], color=color, lw=lw, label=lbl,
                zorder=2, solid_capstyle='butt')

        if mag > 0:
            ux, uy = vx / mag, vy / mag
            px, py = -uy, ux
            hx = ex - ux * hl
            hy = ey - uy * hl
            hw_v = hl * 0.4
            ax.plot([hx + px * hw_v, ex], [hy + py * hw_v, ey],
                    color=color, lw=lw, zorder=3)
            ax.plot([hx - px * hw_v, ex], [hy - py * hw_v, ey],
                    color=color, lw=lw, zorder=3)

    def _draw_double_arrow_2d(ax, ox, oy, vx, vy, color, lbl, arr_len, w, hl, hw):
        """Filled rectangle shaft + filled triangle head (Maple double_arrow)."""
        mag   = np.sqrt(vx**2 + vy**2)
        scale = arr_len / mag if mag > 0 else 1.0
        dx, dy = vx * scale, vy * scale

        # FancyArrow: dx/dy are the full displacement, head_width/head_length/width in data units
        arr = FancyArrow(
            ox, oy, dx, dy,
            width          = w,
            head_width     = hw,
            head_length    = hl,
            length_includes_head = True,
            color          = color,
            zorder         = 2,
        )
        ax.add_patch(arr)

        # Invisible line for legend + combine_plots extents
        ax.plot([ox, ox + dx], [oy, oy + dy], color=color, lw=0.01,
                label=lbl, zorder=1, alpha=0)

    def _draw_harpoon_3d(ax, ox, oy, oz, vx, vy, vz, color, lw, lbl, arr_len, hl):
        """Line shaft + single angled line head — 3D matplotlib."""
        mag   = np.sqrt(vx**2 + vy**2 + vz**2)
        scale = arr_len / mag if mag > 0 else 1.0
        ex = ox + vx * scale
        ey = oy + vy * scale
        ez = oz + vz * scale
        ax.plot([ox, ex], [oy, ey], [oz, ez],
                color=color, lw=lw, label=lbl, zorder=2)
        if mag > 0:
            ux, uy, uz = vx/mag, vy/mag, vz/mag
            perp = np.cross([ux, uy, uz], [0,0,1])
            if np.linalg.norm(perp) < 1e-8:
                perp = np.cross([ux, uy, uz], [0,1,0])
            perp = perp / np.linalg.norm(perp) * hl * 0.5
            hx = ex - ux*hl; hy = ey - uy*hl; hz = ez - uz*hl
            ax.plot([hx + perp[0], ex],
                    [hy + perp[1], ey],
                    [hz + perp[2], ez],
                    color=color, lw=lw, zorder=3)

    def _draw_arrow_3d_shape(ax, ox, oy, oz, vx, vy, vz, color, lw, lbl, arr_len, hl):
        """Line shaft + V head — 3D matplotlib."""
        mag   = np.sqrt(vx**2 + vy**2 + vz**2)
        scale = arr_len / mag if mag > 0 else 1.0
        ex = ox + vx * scale
        ey = oy + vy * scale
        ez = oz + vz * scale
        ax.plot([ox, ex], [oy, ey], [oz, ez],
                color=color, lw=lw, label=lbl, zorder=2)
        if mag > 0:
            ux, uy, uz = vx/mag, vy/mag, vz/mag
            perp = np.cross([ux, uy, uz], [0,0,1])
            if np.linalg.norm(perp) < 1e-8:
                perp = np.cross([ux, uy, uz], [0,1,0])
            perp = perp / np.linalg.norm(perp) * hl * 0.4
            hx = ex - ux*hl; hy = ey - uy*hl; hz = ez - uz*hl
            ax.plot([hx+perp[0], ex], [hy+perp[1], ey], [hz+perp[2], ez],
                    color=color, lw=lw, zorder=3)
            ax.plot([hx-perp[0], ex], [hy-perp[1], ey], [hz-perp[2], ez],
                    color=color, lw=lw, zorder=3)

    def _draw_double_arrow_3d(ax, ox, oy, oz, vx, vy, vz, color, lw, lbl, arr_len, hl):
        """Quiver-based filled arrow — 3D matplotlib approximation of double_arrow."""
        mag   = np.sqrt(vx**2 + vy**2 + vz**2)
        scale = arr_len / mag if mag > 0 else 1.0
        head_frac = hl / arr_len if arr_len > 0 else 0.2
        shaft_end = 1.0 - head_frac
        sx = ox + vx * scale * shaft_end
        sy = oy + vy * scale * shaft_end
        sz = oz + vz * scale * shaft_end
        ax.plot([ox, sx], [oy, sy], [oz, sz],
                color=color, lw=lw * 2, label=lbl, zorder=2)
        ax.quiver(sx, sy, sz,
                  vx * scale * head_frac,
                  vy * scale * head_frac,
                  vz * scale * head_frac,
                  color=color, linewidth=lw,
                  arrow_length_ratio=1.0, zorder=3)

    def _style_2d(ax):
        title_fs, label_fs, tick_fs, leg_fs = _get_font_sizes()
        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.set_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.set_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)
        ax.tick_params(axis='both', labelsize=tick_fs)
        formatter = ticker.ScalarFormatter(useMathText=True)
        formatter.set_scientific(True)
        formatter.set_powerlimits((-2, 2))
        ax.xaxis.set_major_formatter(formatter)
        ax.yaxis.set_major_formatter(formatter)
        if scaling == 'constrained':
            ax.set_aspect('equal', adjustable='box')
        if title: ax.set_title(_smart_label(title), fontsize=title_fs)
        if grid:  ax.grid(True, alpha=0.3)
        if labels:
            ax.legend(fontsize=leg_fs, loc='upper center',
                      bbox_to_anchor=(0.5, -0.15), ncol=2, framealpha=1.0)

    def _style_mpl3d(ax):
        title_fs, label_fs, tick_fs, leg_fs = _get_font_sizes()
        if xlabel: ax.set_xlabel(_smart_label(xlabel), fontsize=label_fs)
        if ylabel: ax.set_ylabel(_smart_label(ylabel), fontsize=label_fs)
        if zlabel: ax.set_zlabel(_smart_label(zlabel), fontsize=label_fs)
        ax.tick_params(axis='both', labelsize=tick_fs)
        if title: ax.set_title(_smart_label(title), fontsize=title_fs)
        if grid:  ax.grid(True)
        if labels:
            ax.legend(fontsize=leg_fs, loc='upper left',
                      bbox_to_anchor=(0.0, 1.0), framealpha=1.0)

    def _render_2d(vecs, origins_np, tips, lims):
        _xlim = xlim or lims[0]
        _ylim = ylim or lims[1]

        fig, ax = plt.subplots(figsize=(width_px / 100, height_px / 100))
        ax.set_xlim(_xlim)
        ax.set_ylim(_ylim)

        for i in range(n):
            ox, oy = origins_np[i]
            vx, vy = vecs[i]
            mag    = np.sqrt(vx**2 + vy**2)
            color  = colors[i] if colors and i < len(colors) else f"C{i}"
            lbl    = _smart_label(labels[i]) if labels and i < len(labels) else None

            arr_len, w, hl, hw = _compute_arrow_dims(mag)

            if shape == 'harpoon':
                _draw_harpoon_2d(ax, ox, oy, vx, vy, color, linewidth, lbl, arr_len, hl)
            elif shape == 'arrow':
                _draw_arrow_2d_shape(ax, ox, oy, vx, vy, color, linewidth, lbl, arr_len, hl)
            elif shape in ('double_arrow', 'cylindrical_arrow'):
                if shape == 'cylindrical_arrow':
                    print("Warning: cylindrical_arrow is 3D only. Using double_arrow for 2D.")
                _draw_double_arrow_2d(ax, ox, oy, vx, vy, color, lbl, arr_len, w, hl, hw)

        _style_2d(ax)

        if show:
            plt.tight_layout()
            plt.show()
        else:
            plt.close()
        return ax

    def _render_mpl3d(vecs, origins_np, tips, lims):
        _xlim = xlim or lims[0]
        _ylim = ylim or lims[1]
        _zlim = zlim or lims[2]

        fig = plt.figure(figsize=(width_px / 100, height_px / 100))
        ax  = fig.add_subplot(111, projection='3d')
        ax.view_init(elev=elev, azim=azim)
        ax.set_xlim(_xlim)
        ax.set_ylim(_ylim)
        ax.set_zlim(_zlim)

        for i in range(n):
            ox, oy, oz = origins_np[i]
            vx, vy, vz = vecs[i]
            mag   = np.sqrt(vx**2 + vy**2 + vz**2)
            color = colors[i] if colors and i < len(colors) else f"C{i}"
            lbl   = _smart_label(labels[i]) if labels and i < len(labels) else None

            arr_len, w, hl, hw = _compute_arrow_dims(mag)

            if shape == 'harpoon':
                _draw_harpoon_3d(ax, ox, oy, oz, vx, vy, vz, color, linewidth, lbl, arr_len, hl)
            elif shape == 'arrow':
                _draw_arrow_3d_shape(ax, ox, oy, oz, vx, vy, vz, color, linewidth, lbl, arr_len, hl)
            else:
                _draw_double_arrow_3d(ax, ox, oy, oz, vx, vy, vz, color, linewidth, lbl, arr_len, hl)

        _style_mpl3d(ax)

        if show:
            plt.tight_layout()
            plt.show()
        else:
            plt.close()
        return ax

    def _render_plotly3d(vecs, origins_np, tips, lims):
        _xlim = xlim or lims[0]
        _ylim = ylim or lims[1]
        _zlim = zlim or lims[2]

        plotly_colors = ['#636EFA','#EF553B','#00CC96','#AB63FA',
                         '#FFA15A','#19D3F3','#FF6692','#B6E880']

        fig = go.Figure()

        for i in range(n):
            ox, oy, oz = origins_np[i]
            vx, vy, vz = vecs[i]
            mag   = np.sqrt(vx**2 + vy**2 + vz**2)
            color = colors[i] if colors and i < len(colors) \
                    else plotly_colors[i % len(plotly_colors)]
            lbl   = labels[i] if labels and i < len(labels) else None
            show_legend = lbl is not None

            arr_len, w, hl, hw = _compute_arrow_dims(mag)
            scale      = arr_len / mag if mag > 0 else 1.0
            head_frac  = hl / arr_len  if arr_len > 0 else 0.2
            shaft_frac = 1.0 - head_frac

            if shape in ('harpoon', 'arrow'):
                # Line shaft + annotate-style head via Scatter3d
                ex = ox + vx * scale
                ey = oy + vy * scale
                ez = oz + vz * scale
                sx = ox + vx * scale * shaft_frac
                sy = oy + vy * scale * shaft_frac
                sz = oz + vz * scale * shaft_frac

                fig.add_trace(go.Scatter3d(
                    x=[ox, ex], y=[oy, ey], z=[oz, ez],
                    mode='lines',
                    line=dict(color=color, width=linewidth * 2),
                    name=lbl or '', showlegend=show_legend,
                    legendgroup=str(i),
                ))
                # Head lines
                if mag > 0:
                    ux, uy, uz = vx/mag, vy/mag, vz/mag
                    perp = np.cross([ux,uy,uz], [0,0,1])
                    if np.linalg.norm(perp) < 1e-8:
                        perp = np.cross([ux,uy,uz], [0,1,0])
                    perp = perp / np.linalg.norm(perp) * hl * 0.4
                    hx = ex - ux*hl; hy = ey - uy*hl; hz = ez - uz*hl
                    lines = [[hx+perp[0], ex], [hy+perp[1], ey], [hz+perp[2], ez]]
                    if shape == 'arrow':
                        lines2 = [[hx-perp[0],ex],[hy-perp[1],ey],[hz-perp[2],ez]]
                    for pts in ([lines] + ([lines2] if shape == 'arrow' else [])):
                        fig.add_trace(go.Scatter3d(
                            x=pts[0], y=pts[1], z=pts[2],
                            mode='lines', line=dict(color=color, width=linewidth*2),
                            showlegend=False, legendgroup=str(i)
                        ))

            else:
                # double_arrow / cylindrical_arrow → Scatter3d shaft + Cone head
                # Calculate absolute tip position
                ex = ox + vx * scale
                ey = oy + vy * scale
                ez = oz + vz * scale
                
                # Calculate where shaft ends
                sx = ox + vx * scale * shaft_frac
                sy = oy + vy * scale * shaft_frac
                sz = oz + vz * scale * shaft_frac

                shaft_width = w * 10 if shape == 'cylindrical_arrow' else linewidth * 3

                # Draw Shaft
                fig.add_trace(go.Scatter3d(
                    x=[ox, sx], y=[oy, sy], z=[oz, sz],
                    mode='lines',
                    line=dict(color=color, width=shaft_width),
                    name=lbl or '', showlegend=show_legend,
                    legendgroup=str(i),
                ))
                
                fig.add_trace(go.Cone(
                    x=[ex], y=[ey], z=[ez],
                    u=[vx], v=[vy], w=[vz],
                    colorscale=[[0, color], [1, color]],
                    showscale=False,
                    sizemode='absolute',
                    sizeref=hl,
                    anchor='tip',
                    name='', showlegend=False,
                    legendgroup=str(i),
                ))

        fig.update_layout(
            title    = title or '',
            autosize = False,
            width    = width_px,
            height   = height_px,
            scene    = dict(
                xaxis_title = xlabel or 'x',
                yaxis_title = ylabel or 'y',
                zaxis_title = zlabel or 'z',
                xaxis       = dict(range=list(_xlim)),
                yaxis       = dict(range=list(_ylim)),
                zaxis       = dict(range=list(_zlim)),
                aspectmode  = 'data' if scaling == 'constrained' else 'cube',
                camera      = dict(eye=dict(x=1.87, y=0.88, z=1.5)),
            ),
            margin = dict(l=10, r=10, t=50, b=10),
            legend = dict(x=0, y=1),
        )

        if show:
            fig.show()
        return fig

    # MAIN EXECUTION

    valid_shapes = ('harpoon', 'arrow', 'double_arrow', 'cylindrical_arrow')
    if shape not in valid_shapes:
        raise ValueError(f"Invalid shape='{shape}'. Choose from: {valid_shapes}")
    if backend not in (None, 'matplotlib', 'plotly'):
        raise ValueError(f"Invalid backend='{backend}'. "
                         f"Choose None (auto), 'matplotlib', or 'plotly'.")

    if not isinstance(vectors, list):
        vectors = [vectors]
    n = len(vectors)

    vecs = [_to_vec(v) for v in vectors]

    dim = len(vecs[0])
    if dim not in (2, 3):
        raise ValueError(f"Vectors must be 2D or 3D, got dimension {dim}.")

    if origins is None:
        origins_np = [np.zeros(dim)] * n
    else:
        if not isinstance(origins, list):
            origins = [origins]
        origins_np = [_to_vec(o) for o in origins]

    tips = [origins_np[i] + vecs[i] for i in range(n)]
    lims = _auto_limits(tips, origins_np)

    # Auto backend and shape defaults
    if backend is None:
        _backend = 'matplotlib' if dim == 2 else 'plotly'
    else:
        _backend = backend

    # Maple default shapes
    _shape = shape
    if _shape == 'double_arrow' and dim == 3 and backend is None:
        _shape = 'cylindrical_arrow'

    if dim == 2:
        return _render_2d(vecs, origins_np, tips, lims)

    if _backend == 'matplotlib':
        return _render_mpl3d(vecs, origins_np, tips, lims)
    else:
        return _render_plotly3d(vecs, origins_np, tips, lims)