import matplotlib.pyplot as plt
from   matplotlib.text import Annotation
from   mpl_toolkits.mplot3d import Axes3D
import plotly.graph_objects as go
import copy
import numpy as np

def combine_plots(plot_list,
                # --- Plot Labels ---
                title=None, xlabel=None, ylabel=None, zlabel=None, labels=None,
                # --- Styling ---
                colors=None, line_styles=None,
                # --- Axis Options ---
                swap_axes=False, center_axes=True, grid=True, scaling=None,
                # --- Limits ---
                xlim=None, ylim=None, zlim=None,
                # --- View angle (matplotlib only) ---
                elev=20,azim=45,
                # --- Display ---
                show=True,
                # --- Figure Size ---
                width=800, height=700):

    """
    Combines multiple plots into a single figure.
    Automatically dispatches by type:
        - matplotlib Axes     → 2D matplotlib handler
        - matplotlib Axes3D   → 3D matplotlib handler
        - plotly Figure       → Plotly handler

    Parameters:
        plot_list    : List of matplotlib Axes / Axes3D or plotly Figure objects
        labels       : Labels for each plot (default=None)
        line_styles  : Line styles or markers to override (default=None) — 2D only
        colors       : Colors for each plot (default=None → keep original)
        swap_axes    : Swap x and y axes — 2D matplotlib only (default=False)
        center_axes  : Move axes to origin — 2D matplotlib only (default=True)
        grid         : Whether to display a grid (default=True)
        scaling      : 'constrained' → equal axis units (default=None)
        xlim         : X-axis limits as (min, max) (default=None)
        ylim         : Y-axis limits as (min, max) (default=None)
        zlim         : Z-axis limits as (min, max) — 3D only (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)
        show         : Whether to display the combined plot (default=True)
        width        : Figure width in pixels (default=800)
        height       : Figure height in pixels (default=700)

    Returns:
        matplotlib Axes   (2D matplotlib input)
        matplotlib Axes3D (3D matplotlib input)
        plotly Figure     (plotly input)
    """

    # Smart Labels
    def _smart_label(lbl):
        if isinstance(lbl, str):
            if any(c in lbl for c in ['_', '^', '\\']):
                return f"${lbl}$"
            return lbl
        return str(lbl)

    # Auto padding — 2D
    def _apply_padding_2d(ax, all_x, all_y):
        if not all_x or not all_y:
            return
        flat_x = np.concatenate([np.atleast_1d(a) for a in all_x])
        flat_y = np.concatenate([np.atleast_1d(a) for a in all_y])
        min_x, max_x = np.nanmin(flat_x), np.nanmax(flat_x)
        min_y, max_y = np.nanmin(flat_y), np.nanmax(flat_y)
        pad_x = (max_x - min_x) * 0.1 if max_x != min_x else 0.5
        pad_y = (max_y - min_y) * 0.1 if max_y != min_y else 0.5
        ax.set_xlim(min_x - pad_x, max_x + pad_x)
        ax.set_ylim(min_y - pad_y, max_y + pad_y)

    # Auto padding — 3D
    def _apply_padding_3d(ax, all_x, all_y, all_z):
        if not all_x or not all_y or not all_z:
            return
        for flat, setter in [
            (np.concatenate([np.atleast_1d(a) for a in all_x]), ax.set_xlim),
            (np.concatenate([np.atleast_1d(a) for a in all_y]), ax.set_ylim),
            (np.concatenate([np.atleast_1d(a) for a in all_z]), ax.set_zlim),
        ]:
            mn, mx = np.nanmin(flat), np.nanmax(flat)
            pad = (mx - mn) * 0.1 if mx != mn else 0.5
            setter(mn - pad, mx + pad)

    # Center axes — 2D only
    def _apply_center_axes(ax):
        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')
        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')

    # Copy 2D arrow annotations
    def _copy_annotations_2d(src_ax, dst_ax, color_override=None):
        for ann in src_ax.get_children():
            if not isinstance(ann, Annotation):
                continue
            if not ann.arrowprops:
                continue
            props = dict(ann.arrowprops)
            if color_override:
                props['color'] = color_override
            try:
                xytext = ann.xyann
            except AttributeError:
                xytext = ann.xy
            dst_ax.annotate('',
                xy         = ann.xy,
                xytext     = xytext,
                xycoords   = ann.xycoords,
                textcoords = getattr(ann, 'anncoords', ann.xycoords),
                arrowprops = props,
                zorder     = ann.get_zorder()
            )

    # 2D Matplotlib
    def _handle_matplotlib_2d(plot_list):
        fig, ax = plt.subplots(figsize=(width / 100, height / 100))

        marker_symbols = ['o', 's', '^', 'x', '*', 'D', 'p', '+', 'v',
                          '<', '>', '1', '2', '3', '4']
        all_plot_x, all_plot_y = [], []

        for i, plot_ax in enumerate(plot_list):
            first_line = True

            # Copy lines
            for line in plot_ax.lines:
                x, y = line.get_data()

                props = {
                    'color'          : line.get_color(),
                    'marker'         : line.get_marker(),
                    'linestyle'      : line.get_linestyle(),
                    'markersize'     : line.get_markersize(),
                    'linewidth'      : line.get_linewidth(),
                    'alpha'          : line.get_alpha(),
                    'zorder'         : line.get_zorder(),
                    'solid_capstyle' : line.get_solid_capstyle(),
                    'solid_joinstyle': line.get_solid_joinstyle()
                }
                if props['marker']    in ['None', '']: props['marker']    = None
                if props['linestyle'] in ['None', '']: props['linestyle'] = 'None'

                if first_line:
                    if line_styles and i < len(line_styles):
                        ns = line_styles[i]
                        if ns in marker_symbols:
                            props['marker'] = ns; props['linestyle'] = 'None'
                        else:
                            props['linestyle'] = ns; props['marker'] = None
                    if colors and i < len(colors):
                        props['color'] = colors[i]

                lbl = labels[i] if first_line and labels and i < len(labels) else None

                if swap_axes:
                    ax.plot(y, x, label=lbl, **props)
                    all_plot_x.append(y); all_plot_y.append(x)
                else:
                    ax.plot(x, y, label=lbl, **props)
                    all_plot_x.append(x); all_plot_y.append(y)

                first_line = False

            # Copy arrow annotations
            _copy_annotations_2d(plot_ax, ax,
                color_override=colors[i] if colors and i < len(colors) else None)

        _apply_padding_2d(ax, all_plot_x, all_plot_y)

        if xlim:  ax.set_xlim(xlim)
        if ylim:  ax.set_ylim(ylim)
        if title: ax.set_title(_smart_label(title))
        if grid:  ax.grid(True)

        if center_axes:
            _apply_center_axes(ax)
        else:
            if xlabel: ax.set_xlabel(_smart_label(xlabel))
            if ylabel: ax.set_ylabel(_smart_label(ylabel))

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

        if labels:
            ax.legend(fontsize=12, loc='upper center',
                      bbox_to_anchor=(0.5, -0.15), ncol=2, framealpha=1.0)

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

    # 3D Matplotlib
    def _handle_matplotlib_3d(plot_list):
        fig = plt.figure(figsize=(width / 100, height / 100))
        ax  = fig.add_subplot(111, projection='3d')
        ax.view_init(elev=elev, azim=azim) 

        all_x, all_y, all_z = [], [], []

        for i, src_ax in enumerate(plot_list):
            color_override = colors[i] if colors and i < len(colors) else None
            first_line     = True

            # Copy line shafts
            for line in src_ax.lines:
                try:
                    xs, ys, zs = line.get_data_3d()
                except Exception:
                    continue

                props = dict(
                    color  = color_override or line.get_color(),
                    lw     = line.get_linewidth(),
                    zorder = line.get_zorder(),
                )
                lbl = None
                if first_line and labels and i < len(labels):
                    lbl = _smart_label(labels[i])
                elif line.get_label() not in ('_nolegend_', ''):
                    lbl = line.get_label()

                ax.plot(xs, ys, zs, label=lbl, **props)
                all_x.append(xs); all_y.append(ys); all_z.append(zs)
                first_line = False

            # Copy surfaces (plot_surface creates a Poly3DCollection).
            # deepcopy is intentionally avoided here: it transitively copies
            # src_ax (including its uninitialised M=None), which causes
            # Poly3DCollection.do_3d_projection() to crash with
            # "NoneType * float" when the new figure is rendered.
            # Instead we reconstruct from the raw vertex / slice data.
            from mpl_toolkits.mplot3d.art3d import Poly3DCollection as _P3C
            for col in src_ax.collections:
                try:
                    if not isinstance(col, _P3C):
                        continue
                    # _vec  : (4, N) homogeneous coords of all vertices
                    # _segslices : list of slice objects, one per face
                    faces = [col._vec[:3, sl].T.tolist()
                             for sl in col._segslices]
                    new_col = _P3C(faces)
                    new_col.set_alpha(col.get_alpha())
                    new_col.set_facecolor(
                        color_override if color_override
                        else col.get_facecolor()
                    )
                    new_col.set_edgecolor(col.get_edgecolor())
                    ax.add_collection3d(new_col)
                except Exception:
                    pass

        _apply_padding_3d(ax, all_x, all_y, all_z)

        if xlim:   ax.set_xlim(xlim)
        if ylim:   ax.set_ylim(ylim)
        if zlim:   ax.set_zlim(zlim)
        if title:  ax.set_title(_smart_label(title))
        if xlabel: ax.set_xlabel(_smart_label(xlabel))
        if ylabel: ax.set_ylabel(_smart_label(ylabel))
        if zlabel: ax.set_zlabel(_smart_label(zlabel))
        if grid:   ax.grid(True)

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

        if labels:
            ax.legend(fontsize=12, loc='upper left',
                      bbox_to_anchor=(0.0, 1.0), framealpha=1.0)

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

    # Plotly
    def _handle_plotly(plot_list):
        combined_fig = go.Figure()

        for i, fig in enumerate(plot_list):
            for trace in fig.data:
                new_trace = copy.deepcopy(trace)
                if colors and i < len(colors):
                    if hasattr(new_trace, 'line'):   new_trace.line.color   = colors[i]
                    if hasattr(new_trace, 'marker'): new_trace.marker.color = colors[i]
                if line_styles and i < len(line_styles):
                    if isinstance(new_trace, go.Scatter) and hasattr(new_trace, 'line'):
                        new_trace.line.dash = line_styles[i]
                if labels and i < len(labels):
                    new_trace.name = labels[i]
                if swap_axes and hasattr(new_trace, 'x') and hasattr(new_trace, 'y'):
                    new_trace.x, new_trace.y = new_trace.y, new_trace.x
                combined_fig.add_trace(new_trace)

        if title:      combined_fig.update_layout(title=title)
        if not center_axes:
            combined_fig.update_layout(xaxis_title=xlabel, yaxis_title=ylabel)
        if xlim:       combined_fig.update_xaxes(range=xlim)
        if ylim:       combined_fig.update_yaxes(range=ylim)
        if grid:       combined_fig.update_layout(xaxis_showgrid=True,
                                                   yaxis_showgrid=True)
        if scaling == 'constrained':
            combined_fig.update_yaxes(scaleanchor="x", scaleratio=1)

        if center_axes:
            combined_fig.update_layout(
                xaxis=dict(zeroline=True, zerolinewidth=1.5,
                           zerolinecolor='black', showline=False,
                           title=xlabel or '', title_standoff=0),
                yaxis=dict(zeroline=True, zerolinewidth=1.5,
                           zerolinecolor='black', showline=False,
                           title=ylabel or '', title_standoff=0)
            )

        combined_fig.update_layout(
            autosize=False, width=width, height=height,
            scene=dict(
                aspectratio=dict(x=1, y=1, z=1),
                camera=dict(eye=dict(x=1.87, y=0.88, z=1.5)),
                xaxis_title=xlabel or '', yaxis_title=ylabel or '',
                zaxis_title=zlabel or '',
                xaxis=dict(range=list(xlim) if xlim else None),
                yaxis=dict(range=list(ylim) if ylim else None),
                zaxis=dict(range=list(zlim) if zlim else None),
            ),
            margin=dict(l=10, r=10, t=50, b=10),
            legend=dict(x=0, y=1),
        )

        if show: combined_fig.show()
        return combined_fig

    # Main Execution
    if not plot_list:
        raise ValueError("plot_list cannot be empty")

    first_plot = plot_list[0]

    if isinstance(first_plot, Axes3D):
        return _handle_matplotlib_3d(plot_list)
    elif hasattr(first_plot, 'lines'):
        return _handle_matplotlib_2d(plot_list)
    elif hasattr(first_plot, 'data'):
        return _handle_plotly(plot_list)
    else:
        raise TypeError(f"Unknown plot type: {type(first_plot)}. "
                        f"Expected matplotlib Axes, Axes3D, or plotly Figure.")