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 ---
                grid=False,
                # --- Limits ---
                xlim=None, ylim=None,
                # --- Display ---
                show=True,
                # --- Figure Size ---
                width=800, height=700):

    """
    Stitches multiple plots end-to-end (extend) or overlays them using internal helpers.

    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 (default=None)
        line_styles : Line styles for each plot (default=None)
        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)
        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)

    # 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

    # Matplotlib Handler
    def _handle_matplotlib(plots):
        fig, ax = plt.subplots(figsize=(width/100, height/100))
        
        agg_xlim = [float('inf'), float('-inf')]
        prev_max_x = None

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

            # 1. Analyze Data Range (for stitching)
            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
            if not has_data:
                shift = 0
            else:
                shift = _get_x_shift(plot_min_x, plot_max_x, prev_max_x, i)

            # 3. Plot Lines
            for line_idx, line in enumerate(lines):
                x_orig, y_orig = line.get_data()
                x, y = x_orig + shift, y_orig
                
                # Style & Color
                c = colors[i] if i < len(colors) else line.get_color()
                
                ls = line_styles[i] if i < len(line_styles) else line.get_linestyle()
                if ls in ['None', '']: ls = ''
                
                mk = line.get_marker()
                if mk in ['None', '']: mk = ''
                
                lbl = labels[i] if (i < len(labels) and line_idx == 0) else None

                ax.plot(x, y, color=c, linestyle=ls, marker=mk, 
                        markersize=line.get_markersize(), linewidth=line.get_linewidth(),
                        label=lbl)

                if len(x) > 0:
                    agg_xlim[0] = min(agg_xlim[0], np.min(x))
                    agg_xlim[1] = max(agg_xlim[1], np.max(x))

            # Update State
            if has_data:
                prev_max_x = plot_max_x + shift

        # Final Config
        final_xlim = xlim if xlim else agg_xlim
        # Use first plot's ylim if not provided
        final_ylim = ylim if ylim is not None else (plots[0].get_ylim() if hasattr(plots[0], 'get_ylim') else None)

        ax.set_xlim(final_xlim)
        if final_ylim: ax.set_ylim(final_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 '')

        ax.set_title(_title)
        ax.set_xlabel(_xlabel)
        ax.set_ylabel(_ylabel)

        if grid: ax.grid(True)
        if any(labels): ax.legend()

        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
            if not has_data:
                shift = 0
            else:
                shift = _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
                
                # Apply Styles
                if i < len(colors):
                    if hasattr(new_trace, 'line'): new_trace.line.color = colors[i]
                    elif hasattr(new_trace, 'marker'): new_trace.marker.color = colors[i]
                
                if i < len(line_styles) and hasattr(new_trace, 'line'):
                    new_trace.line.dash = line_styles[i]
                
                if i < len(labels):
                    new_trace.name = labels[i]
                
                combined_fig.add_trace(new_trace)

            # Update State
            if has_data:
                prev_max_x = plot_max_x + shift

        # Final Layout
        combined_fig.update_layout(
            title=title if title else (layout_defaults.title.text if layout_defaults.title else ''),
            xaxis_title=xlabel if xlabel else (layout_defaults.xaxis.title.text if layout_defaults.xaxis.title else ''),
            yaxis_title=ylabel if ylabel else (layout_defaults.yaxis.title.text if layout_defaults.yaxis.title else '')
        )
        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

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