import matplotlib.pyplot as plt
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, labels=None,
                # --- Styling ---
                colors=None, line_styles=None,
                # --- Axis Options ---
                swap_axes=False, grid=False,
                # --- Limits ---
                xlim=None, ylim=None,
                # --- Display ---
                show=True,
                # --- Figure Size ---
                width=800, height=700):

    """
    Combines multiple Matplotlib Axes or Plotly Figures into a single plot while preserving styles.

    Parameters:
        plot_list  : List of matplotlib Axes or plotly Figure objects to combine
        labels     : Labels for each plot (default=None)
        line_styles: Line styles or markers to override the originals (default=None)
        colors     : Colors for each plot, overriding original colors if provided (default=None)
        swap_axes  : If True, swaps x and y axes for all plots (default=False)
        show       : Whether to display the combined plot (default=True)
        grid       : Whether to display a grid (default=False)
        xlim       : X-axis limits as (min, max) (default=None)
        ylim       : Y-axis limits as (min, max) (default=None)
        title      : Plot title (default=None)
        xlabel     : X-axis label (default=None)
        ylabel     : Y-axis label (default=None)
        width      : width of the Plotly (default=800)
        height     : height of the Plotly (default=700)

    Returns:
        matplotlib Axes or plotly.graph_objects.Figure
            The combined plot object, matching the input plot type
    """
    
    # Apply Padding (Matplotlib)
    def _apply_padding(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)

    # Matplotlib Handler
    def _handle_matplotlib(plot_list):
        # Apply Width and Height (Pixels -> Inches approx / 100)
        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
            for line in plot_ax.lines:
                x, y = line.get_data()

                # --- EXTRACT ALL STYLES ---
                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):
                        new_style = line_styles[i]
                        if new_style in marker_symbols:
                            props['marker'] = new_style
                            props['linestyle'] = 'None'
                        else:
                            props['linestyle'] = new_style
                            props['marker'] = None

                    if colors and i < len(colors):
                        props['color'] = colors[i]
                
                label_to_use = labels[i] if first_line and labels and i < len(labels) else None

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

                first_line = False
        
        _apply_padding(ax, all_plot_x, all_plot_y)

        if xlim: ax.set_xlim(xlim)
        if ylim: ax.set_ylim(ylim)
        if title: ax.set_title(title)
        if xlabel: ax.set_xlabel(xlabel)
        if ylabel: ax.set_ylabel(ylabel)
        if grid: ax.grid(True)
        if labels: ax.legend()
        
        if show: plt.show()
        else: plt.close()
            
        return ax

    # Plotly Handler
    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)
        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)
            
        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))),
            margin=dict(l=10, r=10, t=50, b=10),
            legend=dict(x=0, y=1),
        )

        if show: combined_fig.show()
        return combined_fig

    if not plot_list: raise ValueError("plot_list cannot be empty")
    first_plot = plot_list[0]
    if hasattr(first_plot, 'lines'): return _handle_matplotlib(plot_list)
    elif hasattr(first_plot, 'data'): return _handle_plotly(plot_list)
    else: raise TypeError("Unknown plot object type passed to combine_plots.")