import numpy as np
import sympy as sp
import plotly.graph_objects as go
import plotly.io as pio
import matplotlib.pyplot as plt
from   mpl_toolkits.mplot3d import Axes3D
pio.renderers.default = "notebook"

def plot_3d(exprs, var,
            # --- Backend ---
            backend='plotly',
            # --- Plot Labels ---
            title="3D Plot", xlabel="x", ylabel="y", zlabel="Value", labels=None,
            # --- Styling ---
            colormap=None, show_colorbar=True, opacity=0.9,
            # --- Limits ---
            xlim=None, ylim=None, zlim=None,
            # --- View & Axes Options ---
            axes='boxed', scaling='unconstrained',
            # --- View angle (matplotlib only) ---
            elev=20,azim=45,
            # --- Resolution ---
            resolution=100,
            # --- Display ---
            show=True,
            # --- Figure Size ---
            width=800, height=700):

    """
    Plots 3D surfaces from symbolic expressions using Plotly or Matplotlib.

    Parameters:
        exprs         : SymPy expression or list of expressions
                        - Scalar surface:    z = f(x, y)
                        - Parametric surface: (x(u,v), y(u,v), z(u,v))
        var           : (x_sym, x_range, y_sym, y_range)

        # --- Backend ---
        backend       : 'plotly'     → interactive Plotly surface (default)
                        'matplotlib' → static Matplotlib surface

        # --- Plot Labels ---
        title         : Plot title (default="3D Plot")
        xlabel        : X-axis label (default="x")
        ylabel        : Y-axis label (default="y")
        zlabel        : Z-axis label (default="Value")
        labels        : List of labels for surfaces

        # --- Styling ---
        colormap      : Colormap specification:
                        - None          : default (Viridis/viridis) for all surfaces
                        - String        : apply same colormap to ALL surfaces
                        - List          : apply colormap[i] to surface i
        show_colorbar : Whether to display colorbar (default=True)
        opacity       : Surface opacity — float or list per surface (default=0.9)

        # --- 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) (default=None → auto)

        # --- View & Axes Options ---
        axes          : Plotly only — "boxed", "normal", "frame", or "none" (default="boxed")
        scaling       : "unconstrained" or "constrained" (default="unconstrained")

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

        # --- Resolution ---
        resolution    : Grid resolution per axis (default=100)

        # --- Display ---
        show          : Whether to display the plot (default=True)

        # --- Figure Size ---
        width         : Figure width in pixels (default=800)
        height        : Figure height in pixels (default=700)

    Returns:
        plotly.graph_objects.Figure  (backend='plotly')
        matplotlib Axes3D            (backend='matplotlib')
    """

    # Normalise colormap input
    def _get_colorscale_plotly(cmap_input):
        SINGLE_COLORS = {
            'red','blue','green','yellow','orange','purple','pink',
            'cyan','magenta','brown','gray','grey','black','white',
            'navy','teal','lime','olive','maroon','aqua','fuchsia',
            'silver','gold','indigo','violet','crimson','coral'
        }
        if cmap_input is None:
            return "Viridis"
        s = str(cmap_input).lower()
        if s in SINGLE_COLORS or s.startswith('#') or s.startswith('rgb'):
            return [[0.0, cmap_input], [1.0, cmap_input]]
        return cmap_input

    # Compute parametric surface arrays
    def _compute_parametric(expr_tuple, u_sym, v_sym, U, V):
        ex, ey, ez = [sp.sympify(e) for e in expr_tuple]
        fx = sp.lambdify((u_sym, v_sym), ex, modules='numpy')
        fy = sp.lambdify((u_sym, v_sym), ey, modules='numpy')
        fz = sp.lambdify((u_sym, v_sym), ez, modules='numpy')
        X = fx(U, V); Y = fy(U, V); Z = fz(U, V)
        if np.isscalar(X): X = np.full_like(U, X, dtype=float)
        if np.isscalar(Y): Y = np.full_like(U, Y, dtype=float)
        if np.isscalar(Z): Z = np.full_like(U, Z, dtype=float)
        return X, Y, Z

    # Compute scalar surface arrays
    def _compute_scalar(expr_scalar, u_sym, v_sym, U, V):
        ez  = sp.sympify(expr_scalar)
        fz  = sp.lambdify((u_sym, v_sym), ez, modules='numpy')
        Z   = fz(U, V)
        if np.isscalar(Z): Z = np.full_like(U, Z, dtype=float)
        return U, V, Z

    # Build grids and surface list
    def _prepare_surfaces():
        u_min, u_max = float(var[1][0]), float(var[1][1])
        v_min, v_max = float(var[3][0]), float(var[3][1])
        u_vals = np.linspace(u_min, u_max, resolution)
        v_vals = np.linspace(v_min, v_max, resolution)
        U, V   = np.meshgrid(u_vals, v_vals)

        if isinstance(exprs, tuple):   surface_list = [exprs]
        elif isinstance(exprs, list):  surface_list = exprs
        else:                          surface_list = [exprs]

        # Colormap per surface
        if colormap is None:                       cmaps = [None]     * len(surface_list)
        elif isinstance(colormap, str):            cmaps = [colormap] * len(surface_list)
        elif isinstance(colormap, (list, tuple)):  cmaps = list(colormap)
        else:                                      cmaps = [colormap] * len(surface_list)

        # Opacity per surface
        if isinstance(opacity, (list, tuple)):  opacities = list(opacity)
        else:                                   opacities = [opacity] * len(surface_list)

        surfaces = []
        for i, item in enumerate(surface_list):
            lbl   = labels[i] if labels and i < len(labels) else f"Surface {i+1}"
            cmap  = cmaps[i]     if i < len(cmaps)     else None
            alpha = opacities[i] if i < len(opacities) else opacities[-1]

            if isinstance(item, tuple) and len(item) == 3:
                X, Y, Z = _compute_parametric(item, var[0], var[2], U, V)
            else:
                X, Y, Z = _compute_scalar(item, var[0], var[2], U, V)

            surfaces.append(dict(X=X, Y=Y, Z=Z, label=lbl, cmap=cmap, alpha=alpha))

        return surfaces

    # Plotly
    def _render_plotly(surfaces):
        fig = go.Figure()

        for s in surfaces:
            scale = _get_colorscale_plotly(s['cmap'])
            fig.add_trace(go.Surface(
                x=s['X'], y=s['Y'], z=s['Z'],
                name=s['label'],
                colorscale=scale,
                showscale=show_colorbar,
                opacity=s['alpha'],
                showlegend=True
            ))

        scene_aspect = dict(aspectmode='data') if scaling == 'constrained' \
                       else dict(aspectmode='cube')
        show_ax = axes != 'none'
        ax_cfg  = dict(visible=show_ax, showgrid=show_ax,
                       showbackground=show_ax, zeroline=show_ax)

        fig.update_layout(
            title=title,
            autosize=False, width=width, height=height,
            scene=dict(
                xaxis_title=xlabel, yaxis_title=ylabel, zaxis_title=zlabel,
                xaxis=dict(range=list(xlim) if xlim else None, **ax_cfg),
                yaxis=dict(range=list(ylim) if ylim else None, **ax_cfg),
                zaxis=dict(range=list(zlim) if zlim else None, **ax_cfg),
                **scene_aspect,
                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

    # Matplotlib
    def _render_matplotlib(surfaces):
        fig_mpl = plt.figure(figsize=(width / 100, height / 100))
        ax      = fig_mpl.add_subplot(111, projection='3d')
        ax.view_init(elev=elev, azim=azim)

        for s in surfaces:
            # Colormap: None → viridis, single string → try as mpl cmap
            cmap_mpl = s['cmap'] if s['cmap'] is not None else 'viridis'

            ax.plot_surface(
                s['X'], s['Y'], s['Z'],
                cmap    = cmap_mpl,
                alpha   = s['alpha'],
                label   = s['label'],
                zorder  = 2
            )

        if show_colorbar and len(surfaces) == 1:
            mappable = plt.cm.ScalarMappable(
                cmap=surfaces[0]['cmap'] or 'viridis')
            mappable.set_array(surfaces[0]['Z'])
            fig_mpl.colorbar(mappable, ax=ax, shrink=0.5)

        ax.set_title(title)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.set_zlabel(zlabel)

        if xlim: ax.set_xlim(xlim)
        if ylim: ax.set_ylim(ylim)
        if zlim: ax.set_zlim(zlim)

        if scaling == 'constrained':
            ax.set_box_aspect([1, 1, 1])

        if labels:
            ax.legend()

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

        return ax

    # Main Execution
    if not isinstance(var, tuple) or len(var) != 4:
        raise ValueError("`var` must be (u_sym, u_range, v_sym, v_range)")
    if backend not in ('plotly', 'matplotlib'):
        raise ValueError(f"Invalid backend='{backend}'. Choose 'plotly' or 'matplotlib'.")

    surfaces = _prepare_surfaces()

    if backend == 'plotly':
        return _render_plotly(surfaces)
    else:
        return _render_matplotlib(surfaces)