import matplotlib.pyplot as plt
import plotly.graph_objects as go
import copy
import numpy as np

def extend_plot(plot_list, dx=0, stitch=True,
                # --- Plot Labels ---
                title=None, xlabel=None, ylabel=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,
                # --- Display ---
                show=True,
                # --- Figure Size ---
                width=800, height=700):

    """
    Stitches multiple plots end-to-end or overlays them, with full style preservation.

    Parameters:
        plot_list    : List of matplotlib Axes or plotly Figure objects
        dx           : Horizontal offset between consecutive plots (default=0)
        stitch       : If True, appends plots horizontally. If False, overlays them (default=True)
        title        : Plot title (default=None)
        xlabel       : X-axis label (default=None)
        ylabel       : Y-axis label (default=None)
        labels       : Labels for each segment (default=None)
        colors       : Colors for each plot, overriding originals if provided (default=None)
        line_styles  : Line styles or markers to override originals (default=None)
        swap_axes    : If True, swaps x and y axes for all plots (default=False)
        center_axes  : If True, moves x and y axes to pass through the origin (default=False)
        grid         : Whether to display a grid (default=False)
        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)
        show         : Whether to display the plot (default=True)
        width        : Plot width in pixels (default=800)
        height       : Plot height in pixels (default=700)

    Returns:
        matplotlib Axes or plotly.graph_objects.Figure
    """

    if not plot_list:
        raise ValueError("Plot list cannot be empty")

    # Default Handling for Mutable Arguments
    if colors      is None: colors      = []
    if line_styles is None: line_styles = []
    if labels      is None: labels      = []
    if isinstance(colors, str): colors  = [colors] * len(plot_list)

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

    # 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)

    # Calculate X-Shift
    def _get_x_shift(curr_min, curr_max, prev_max, index):
        if not stitch or index == 0:
            return 0
        return (prev_max - curr_min) + dx

    # 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(np.asarray(a, dtype=float)) for a in all_x])
        flat_y = np.concatenate([np.atleast_1d(np.asarray(a, dtype=float)) 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)

    # Apply Center Axes (Matplotlib)
    def _apply_center_axes(ax, xlabel_str, ylabel_str):
        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_str:
            ax.set_xlabel('')
            ax.annotate(_smart_label(xlabel_str),
                        xy=(1, 0), xycoords=('axes fraction', 'data'),
                        xytext=(10, 0), textcoords='offset points',
                        ha='left', va='center')
        if ylabel_str:
            ax.set_ylabel('')
            ax.annotate(_smart_label(ylabel_str),
                        xy=(0, 1), xycoords=('data', 'axes fraction'),
                        xytext=(0, 10), textcoords='offset points',
                        ha='center', va='bottom')

    # Matplotlib Handler
    def _handle_matplotlib(plots):
        fig, ax = plt.subplots(figsize=(width / 100, height / 100))

        all_plot_x = []
        all_plot_y = []
        prev_max_x = None

        for i, p in enumerate(plots):
            if hasattr(p, 'lines'):       lines = p.lines
            elif hasattr(p, 'get_lines'): lines = p.get_lines()
            else:                         lines = []

            # 1. Analyze Data Range
            plot_min_x, plot_max_x = float('inf'), float('-inf')
            has_data = False

            for line in lines:
                x_orig, _ = line.get_data()
                if len(x_orig) > 0:
                    has_data = True
                    plot_min_x = min(plot_min_x, np.min(x_orig))
                    plot_max_x = max(plot_max_x, np.max(x_orig))

            # 2. Calculate Shift
            shift = 0 if not has_data else _get_x_shift(plot_min_x, plot_max_x, prev_max_x, i)

            # 3. Plot Lines — full style fidelity
            first_line = True
            for line in lines:
                x_orig, y_orig = line.get_data()
                x, y = x_orig + shift, y_orig

                # Extract ALL style properties
                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'

                # Override styles on first line only
                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]

                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

            if has_data:
                prev_max_x = plot_max_x + shift

        # Apply Auto-Padding
        _apply_padding(ax, all_plot_x, all_plot_y)

        # Override with manual limits
        if xlim: ax.set_xlim(xlim)
        if ylim: ax.set_ylim(ylim)

        # Labels (Prioritize arg -> source -> empty)
        _title  = title  if title  else (plots[0].get_title()  if hasattr(plots[0], 'get_title')  else '')
        _xlabel = xlabel if xlabel else (plots[0].get_xlabel() if hasattr(plots[0], 'get_xlabel') else '')
        _ylabel = ylabel if ylabel else (plots[0].get_ylabel() if hasattr(plots[0], 'get_ylabel') else '')

        if _title: ax.set_title(_smart_label(_title))
        if grid:   ax.grid(True)
        if labels and any(labels): ax.legend()

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

        # Center Axes overrides normal label placement
        if center_axes:
            _apply_center_axes(ax, _xlabel, _ylabel)
        else:
            if _xlabel: ax.set_xlabel(_smart_label(_xlabel))
            if _ylabel: ax.set_ylabel(_smart_label(_ylabel))

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

    # Plotly Handler
    def _handle_plotly(plots):
        combined_fig    = go.Figure()
        prev_max_x      = None
        layout_defaults = plots[0].layout

        for i, fig in enumerate(plots):
            # 1. Analyze Data Range
            plot_min_x, plot_max_x = float('inf'), float('-inf')
            has_data = False

            for trace in fig.data:
                if hasattr(trace, 'x') and trace.x is not None:
                    x_data = np.array(trace.x)
                    if len(x_data) > 0:
                        has_data = True
                        plot_min_x = min(plot_min_x, np.min(x_data))
                        plot_max_x = max(plot_max_x, np.max(x_data))

            # 2. Calculate Shift
            shift = 0 if not has_data else _get_x_shift(plot_min_x, plot_max_x, prev_max_x, i)

            # 3. Add Traces
            for trace in fig.data:
                new_trace = copy.deepcopy(trace)

                # Apply Shift
                if hasattr(new_trace, 'x') and new_trace.x is not None:
                    new_trace.x = np.array(new_trace.x) + shift

                # Swap Axes
                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

                # Override Color
                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]

                # Override Line Style
                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]

                # Override Label
                if labels and i < len(labels):
                    new_trace.name = labels[i]

                combined_fig.add_trace(new_trace)

            if has_data:
                prev_max_x = plot_max_x + shift

        # Final Layout
        _title  = title  if title  else (layout_defaults.title.text      if layout_defaults.title      else '')
        _xlabel = xlabel if xlabel else (layout_defaults.xaxis.title.text if layout_defaults.xaxis.title else '')
        _ylabel = ylabel if ylabel else (layout_defaults.yaxis.title.text if layout_defaults.yaxis.title else '')

        if _title: combined_fig.update_layout(title=_title)

        if center_axes:
            combined_fig.update_layout(
                xaxis=dict(zeroline=True, zerolinewidth=1.5, zerolinecolor='black',
                           showline=False, title=_xlabel, title_standoff=0),
                yaxis=dict(zeroline=True, zerolinewidth=1.5, zerolinecolor='black',
                           showline=False, title=_ylabel, title_standoff=0)
            )
        else:
            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)

        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

    # Main Execution
    first_plot = plot_list[0]
    if   hasattr(first_plot, 'lines') or hasattr(first_plot, 'get_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 extend_plot.")