import numpy as np
import sympy as sp
import plotly.graph_objects as go
import plotly.io as pio
pio.renderers.default = "notebook"

def plot_3d(exprs, var,
            # --- 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',
            # --- Resolution ---
            resolution=100,
            # --- Display ---
            show=True,
            # --- Figure Size ---
            width=800, height=700):
    
    """
    Plots 3D surfaces from symbolic expressions using Plotly with internal helpers

    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)
        title         : Plot title (default="3D Plot")
        labels        : List of labels for surfaces
        colormap      : Colormap specification:
                        - None: default (Viridis) for all surfaces
                        - String: apply same colormap to ALL surfaces
                        - List: apply colormap[i] to surface i (None for remaining)
        show_colorbar : Whether to display colorbar (default=True)
        resolution    : Grid resolution per axis (default=100)
        show          : Whether to display the plot (default=True)
        width         : Figure width in pixels
        height        : Figure height in pixels
        axes          : "boxed", "normal", "frame", or "none" (default="boxed")
        scaling       : "unconstrained" or "constrained" (default="unconstrained")
        opacity       : Surface opacity
                        - Float          : same opacity for all surfaces
                        - List or tuple  : opacity per surface

    Returns:
        plotly.graph_objects.Figure
    """

    # Color Processing
    def _get_colorscale(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'
        }
        
        str_input = str(cmap_input).lower()
        if cmap_input is None: return "Viridis"
            
        is_single = (str_input in SINGLE_COLORS or 
                     str_input.startswith('#') or 
                     str_input.startswith('rgb'))
        
        if is_single: return [[0.0, cmap_input], [1.0, cmap_input]]
        return cmap_input

    # Parametric Logic
    def _compute_parametric(expr_tuple, u_sym, v_sym, u_grid, v_grid):
        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_vals = fx(u_grid, v_grid)
        Y_vals = fy(u_grid, v_grid)
        Z_vals = fz(u_grid, v_grid)
        
        if np.isscalar(X_vals): X_vals = np.full_like(u_grid, X_vals, dtype=float)
        if np.isscalar(Y_vals): Y_vals = np.full_like(u_grid, Y_vals, dtype=float)
        if np.isscalar(Z_vals): Z_vals = np.full_like(u_grid, Z_vals, dtype=float)
        return X_vals, Y_vals, Z_vals

    # Scalar Logic
    def _compute_scalar(expr_scalar, u_sym, v_sym, u_grid, v_grid):
        ez = sp.sympify(expr_scalar)
        fz = sp.lambdify((u_sym, v_sym), ez, modules="numpy")
        Z_vals = fz(u_grid, v_grid)
        
        if np.isscalar(Z_vals): Z_vals = np.full_like(u_grid, Z_vals, dtype=float)
        return u_grid, v_grid, Z_vals

    # Build Figure
    def _build_figure(u_range, v_range):
        fig = go.Figure()
        
        try:
            u_min, u_max = float(u_range[0]), float(u_range[1])
            v_min, v_max = float(v_range[0]), float(v_range[1])
        except TypeError:
            raise ValueError("Limits must be numeric.")
        
        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)

        # 1. Normalize Input to List
        if isinstance(exprs, tuple): surface_list = [exprs] 
        elif isinstance(exprs, list): surface_list = exprs
        else: surface_list = [exprs]
        
        # Colormap Distribution
        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 Distribution
        if isinstance(opacity, (list, tuple)):
            opacities = list(opacity)
        else:
            opacities = [opacity] * len(surface_list)

        # Loop over surfaces
        for i, item in enumerate(surface_list):
            lbl = labels[i] if labels and i < len(labels) else f"Surface {i + 1}"
            
            # Get specific color and opacity for this iteration
            cmap_raw = cmaps[i] if i < len(cmaps) else None
            curr_opacity = opacities[i] if i < len(opacities) else opacities[-1]
            
            scale = _get_colorscale(cmap_raw)

            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)
            
            fig.add_trace(go.Surface(
                x=X, y=Y, z=Z,
                name=lbl,
                colorscale=scale,
                showscale=show_colorbar,
                opacity=curr_opacity,
                showlegend=True
            ))
        return fig

    # Main Execution
    if not isinstance(var, tuple) or len(var) != 4:
        raise ValueError("`var` must be (u_sym, u_range, v_sym, v_range)")

    final_fig = _build_figure(var[1], var[3])

    # Layout
    scene_aspect = dict(aspectmode='data') if scaling == 'constrained' else dict(aspectmode='cube')
    show_axes = False if axes == 'none' else True
    axis_config = dict(visible=show_axes, showgrid=show_axes, showbackground=show_axes, zeroline=show_axes)

    final_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, **axis_config),
            yaxis=dict(range=list(ylim) if ylim else None, **axis_config),
            zaxis=dict(range=list(zlim) if zlim else None, **axis_config),
            **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:
        final_fig.show()

    return final_fig