import numpy as np
import sympy as sp
import matplotlib.pyplot as plt

def plot_param_grid(surface_exprs, var,
                    # --- Plot Labels ---
                    title="Parametric Surface Grid",
                    xlabel="X", ylabel="Y", labels=None,
                    # --- Styling ---
                    grid_color='blue', points_color='red', cage_color='red',
                    points_label='Control Points',
                    # --- Grid / Cage Options ---
                    plot_cage=True, grid_lines=11,
                    # --- Limits ---
                    xlim=None, ylim=None,
                    # --- Geometry / Control Points ---
                    control_points=None, cp_dims=None,
                    # --- Display ---
                    show=True,
                    # --- Figure Size ---
                    width=800, height=600):

    """
    Plots a 2D parametric surface's isoparametric grid and its control points using internal helpers

    Parameters:
        surface_exprs  : List/tuple of two symbolic expressions [X(u,v), Y(u,v)]
        var            : Tuple (u_sym, u_range, v_sym, v_range)
        control_points : Optional Matrix or array of control point coordinates (N x 2)
        cp_dims        : Tuple (num_rows, num_cols) describing grid layout of control points
        plot_cage      : If True, connects control points to form a cage
        grid_lines     : Number of isoparametric lines to draw in each direction
        width          : Plot width in pixels (default=800)
        height         : Plot height in pixels (default=600)

    Returns:
        fig, ax        : The matplotlib Figure and Axes objects
    """

    # Process & Plot Control Points
    def _plot_control_geometry(ax, cp_data, dims):
        # Convert SymPy Matrix to NumPy if needed
        if isinstance(cp_data, sp.Matrix):
            cp_array = np.array(cp_data.tolist()).astype(np.float64)
        else:
            cp_array = np.array(cp_data)

        # Plot Points
        ax.plot(cp_array[:, 0], cp_array[:, 1], 'o', 
                color=points_color, markersize=8, label=points_label, zorder=5)

        # Plot Cage (Grid connections)
        if plot_cage:
            if dims and len(dims) == 2:
                rows, cols = dims
                # Reshape linear array to (rows, cols, 2)
                # Ensure the total size matches
                if cp_array.shape[0] == rows * cols:
                    grid = cp_array.reshape(rows, cols, 2)
                    
                    # Plot Rows
                    for i in range(rows):
                        ax.plot(grid[i, :, 0], grid[i, :, 1], '-', 
                                color=cage_color, alpha=0.5, linewidth=1.0)
                    # Plot Cols
                    for j in range(cols):
                        ax.plot(grid[:, j, 0], grid[:, j, 1], '-', 
                                color=cage_color, alpha=0.5, linewidth=1.0)
                else:
                    print(f"Warning: Control points size {cp_array.shape[0]} does not match dims {dims} ({rows*cols}). Cage skipped.")
            else:
                print("Warning: `cp_dims` (rows, cols) must be provided to plot the control cage.")

    # Evaluate & Plot Isoparametric Grid
    def _plot_isoparametric_grid(ax, exprs, u_s, u_r, v_s, v_r):
        # Compile functions for fast numerical evaluation
        fx = sp.lambdify((u_s, v_s), exprs[0], modules='numpy')
        fy = sp.lambdify((u_s, v_s), exprs[1], modules='numpy')

        u_min, u_max = float(u_r[0]), float(u_r[1])
        v_min, v_max = float(v_r[0]), float(v_r[1])

        # Resolution for the curves themselves
        curve_res = 100 

        # 1. Constant V lines (vary U)
        u_vals = np.linspace(u_min, u_max, curve_res)
        v_levels = np.linspace(v_min, v_max, grid_lines)
        
        for v_k in v_levels:
            # Broadcast v_k to match u_vals shape
            v_const = np.full_like(u_vals, v_k)
            x_out = fx(u_vals, v_const)
            y_out = fy(u_vals, v_const)
            
            # Handle scalar output (constant functions)
            if np.isscalar(x_out): x_out = np.full_like(u_vals, x_out)
            if np.isscalar(y_out): y_out = np.full_like(u_vals, y_out)
                
            ax.plot(x_out, y_out, color=grid_color, linewidth=1.5, alpha=0.8)

        # 2. Constant U lines (vary V)
        v_vals = np.linspace(v_min, v_max, curve_res)
        u_levels = np.linspace(u_min, u_max, grid_lines)

        for u_k in u_levels:
            u_const = np.full_like(v_vals, u_k)
            x_out = fx(u_const, v_vals)
            y_out = fy(u_const, v_vals)

            if np.isscalar(x_out): x_out = np.full_like(v_vals, x_out)
            if np.isscalar(y_out): y_out = np.full_like(v_vals, y_out)

            ax.plot(x_out, y_out, color=grid_color, linewidth=1.5, alpha=0.8)

    # Main Execution
    
    # Validation
    if not isinstance(surface_exprs, (list, tuple)) or len(surface_exprs) != 2:
        raise ValueError("`surface_exprs` must be [X(u,v), Y(u,v)].")
    if not isinstance(var, tuple) or len(var) != 4:
        raise ValueError("`var` must be (u_sym, u_range, v_sym, v_range)")

    # Setup Figure
    # Convert pixels to inches for Matplotlib (assuming ~100 dpi)
    fig, ax = plt.subplots(figsize=(width/100, height/100))

    # 1. Plot Control Geometry
    if control_points is not None:
        _plot_control_geometry(ax, control_points, cp_dims)

    # 2. Plot Parametric Grid
    _plot_isoparametric_grid(ax, surface_exprs, var[0], var[1], var[2], var[3])

    # 3. Layout & Styling
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    
    if xlim: ax.set_xlim(xlim)
    if ylim: ax.set_ylim(ylim)
    
    ax.grid(True, linestyle=':', alpha=0.6)
    ax.set_aspect('equal', adjustable='datalim')

    # Handle Legend
    handles, lbls = ax.get_legend_handles_labels()
    if lbls:
        # Deduplicate legend entries
        by_label = dict(zip(lbls, handles))
        ax.legend(by_label.values(), by_label.keys())

    if show:
        plt.show()

    return fig, ax