import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from sympy import latex, Matrix
import matplotlib.ticker as ticker

def plot_2d(exprs, var=None,
            # --- Plot Labels ---
            title=None, xlabel=None, ylabel=None, labels=None,
            # --- Styling ---
            colors=None, line_styles=None, linewidth=2, scaling=None, grid=True,
            solid_capstyle='butt',
            # --- Typography ---
            fontsize=14, title_size=None, label_size=None,
            tick_size=None, legend_size=None,
            # --- Axis Options ---
            swap_axes=False, center_axes=True,
            # --- Limits ---
            xlim=None, ylim=None,
            # --- Resolution ---
            resolution=400,
            # --- Display settings ---
            show=True, width=800, height=600,
            # --- Break Control ---
            _break_=None):

    r"""
    Plots 2D curves from symbolic expressions or numeric datasets using Matplotlib.

    Parameters:
        exprs        : Expression, tuple, or list of expressions to plot
                        1. (x, y)               → point plot if numeric vectors, parametric if symbolic
                        2. (x, y, (t, min, max)) → parametric with custom range
                        3. f(x)                  → standard y = f(x) plot
        var          : Symbol or (symbol, (min, max)) defining the variable and range.
                       Optional when plotting numeric vector pairs.
        labels       : Labels for each curve (default=None)
        line_styles  : Line styles or markers for each curve (default=None)
        linewidth    : Line width for curves (default=2)
        colors       : Colors for each curve (default=None)
        title        : Plot title (default=None)
        xlabel       : X-axis label (default=None)
        ylabel       : Y-axis label (default=None)
        xlim         : X-axis limits as (min, max) (default=None)
        ylim         : Y-axis limits as (min, max) (default=None)
        resolution   : Number of points used to evaluate symbolic expressions (default=400)
        show         : Whether to display the plot (default=True)
        _break_      : List of expressions for which discontinuities should be visually broken (default=None)
        fontsize     : Base font size for plot text (default=14)
        title_size   : Font size for the title (default=None → fontsize+2)
        label_size   : Font size for axis labels (default=None → fontsize)
        tick_size    : Font size for tick labels (default=None → fontsize-2)
        legend_size  : Font size for legend text (default=None → tick_size)
        swap_axes    : If True, swaps x and y axes (default=False)
        center_axes  : If True, moves x and y axes to pass through the origin (default=False)
        scaling      : 'constrained' → equal axis units (default=None)
        grid         : Whether to display a grid (default=True)
        solid_capstyle : 'butt', 'round', or 'projecting' (default='butt')

    Returns:
        matplotlib Axes
    """

    # Break visual discontinuities
    def _break_kinks(y_data):
        dy = np.abs(np.diff(y_data))
        max_dy = np.nanmax(dy)
        if max_dy == 0 or np.isnan(max_dy):
            return y_data
        threshold = max_dy * 0.5
        y_data = y_data.copy()
        for idx in np.where(dy > threshold)[0]:
            y_data[idx] = np.nan
            y_data[idx + 1] = np.nan
        return y_data

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

    # Convert to numpy array if SymPy matrix
    def _to_numpy(data):
        if isinstance(data, sp.Matrix):
            return np.array(data).astype(np.float64).flatten()
        if isinstance(data, np.ndarray):
            return data.astype(np.float64).flatten()
        return data

    # Check if data is a numeric array (not symbolic)
    def _is_numeric(data):
        return isinstance(data, (sp.Matrix, np.ndarray)) and not isinstance(data, sp.Basic)

    # Handle Parametric Logic
    def _get_parametric_data(expr, x_sym, x_vals_sample):
        # Case A: Custom Range (x_expr, y_expr, (t, min, max))
        if len(expr) == 3:
            x_e, y_e, range_info = expr
            t_sym, t_min, t_max = range_info

            t_vals = np.linspace(float(t_min), float(t_max), resolution)
            f_x = sp.lambdify(t_sym, x_e, modules='numpy')
            f_y = sp.lambdify(t_sym, y_e, modules='numpy')

            x_res = f_x(t_vals)
            y_res = f_y(t_vals)

            if not hasattr(x_res, '__len__'): x_res = np.full_like(t_vals, x_res)
            if not hasattr(y_res, '__len__'): y_res = np.full_like(t_vals, y_res)

            return np.array(x_res).flatten(), np.array(y_res).flatten(), False

        # Case B: Pair (x, y)
        elif len(expr) == 2:
            x_data, y_data = expr

            # Convert SymPy scalar constants → numpy so they hit the numeric branch
            if isinstance(x_data, sp.Basic) and not x_data.free_symbols:
                x_data = np.array([float(x_data)])
            if isinstance(y_data, sp.Basic) and not y_data.free_symbols:
                y_data = np.array([float(y_data)])

            # --- Numeric vector pair → plot directly ---
            if _is_numeric(x_data) and _is_numeric(y_data):
                x_data = _to_numpy(x_data)
                y_data = _to_numpy(y_data)
                if expr in _break_:
                    y_data = _break_kinks(y_data)
                return x_data, y_data, True   # True = is_numeric_pair

            # --- Symbolic parametric ---
            if isinstance(x_data, sp.Basic) or isinstance(y_data, sp.Basic):
                f_x = sp.lambdify(x_sym, x_data, modules='numpy')
                f_y = sp.lambdify(x_sym, y_data, modules='numpy')
                x_data = f_x(x_vals_sample)
                y_data = f_y(x_vals_sample)

                if hasattr(x_data, '__len__') and not hasattr(y_data, '__len__'):
                    y_data = np.full_like(x_data, y_data)
                elif not hasattr(x_data, '__len__') and hasattr(y_data, '__len__'):
                    x_data = np.full_like(y_data, x_data)

            if expr in _break_:
                y_data = _break_kinks(y_data)

            return x_data, y_data, False

        return None, None, False

    # Handle Standard y=f(x) Logic
    def _get_standard_data(expr, x_sym, x_vals_sample):
        original_expr = expr
        expr = sp.sympify(expr)

        if not expr.has(x_sym):
            y_vals = np.full_like(x_vals_sample, float(expr))
        else:
            f = sp.lambdify(x_sym, expr, modules='numpy')
            y_vals = np.array(f(x_vals_sample)).flatten()

        if original_expr in _break_:
            y_vals = _break_kinks(y_vals)

        return x_vals_sample, y_vals

    # Apply Padding
    def _apply_padding(ax, all_x, all_y):
        if not all_x or not all_y:
            return

        flat_x = np.concatenate([np.asarray(a, dtype=float).ravel() for a in all_x])
        flat_y = np.concatenate([np.asarray(a, dtype=float).ravel() 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.2 if max_x != min_x else 0.5
        pad_y = (max_y - min_y) * 0.2 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)

    # MAIN EXECUTION

    # Font/Size Config
    title_size  = title_size  or fontsize + 2
    label_size  = label_size  or fontsize
    tick_size   = tick_size   or fontsize - 2
    legend_size = legend_size or tick_size

    # Input Normalization
    if not isinstance(exprs, list):
        exprs = [exprs]
    if _break_ is None:
        _break_ = []

    # Parse var (optional)
    x_sym         = None
    x_vals_sample = None
    if var is not None:
        if isinstance(var, tuple):
            x_sym, x_range = var
        else:
            x_sym   = var
            x_range = (-1, 1)
        x_vals_sample = np.linspace(float(x_range[0]), float(x_range[1]), resolution)

    # Setup Plot
    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 = []

    # Iterate Expressions
    for i, expr in enumerate(exprs):
        style     = line_styles[i] if line_styles and i < len(line_styles) else 'solid'
        color     = colors[i]      if colors      and i < len(colors)      else None
        raw_label = labels[i]      if labels      and i < len(labels)      else None
        label     = _smart_label(raw_label) if raw_label is not None else None

        marker    = None
        is_numeric_pair = False

        if isinstance(expr, Matrix):
            expr = np.array(expr).astype(np.float64).flatten()

        final_x = None
        final_y = None

        # A. Parametric / numeric pair (tuple or list)
        if isinstance(expr, (tuple, list)):
            final_x, final_y, is_numeric_pair = _get_parametric_data(expr, x_sym, x_vals_sample)

            # Auto point-plot for numeric vector pairs when no style specified
            if is_numeric_pair and style == 'solid':
                marker = 'o'
                style  = ''

        # B. Standard y = f(x)
        else:
            if x_sym is None or x_vals_sample is None:
                raise ValueError("var must be provided for symbolic expressions.")
            final_x, final_y = _get_standard_data(expr, x_sym, x_vals_sample)

        # Handle explicit marker styles
        if style in marker_symbols:
            marker = style
            style  = ''

        # Plot
        if final_x is not None and final_y is not None:
            if swap_axes:
                ax.plot(final_y, final_x, label=label, linestyle=style,
                        color=color, marker=marker, linewidth=linewidth,
                        solid_capstyle=solid_capstyle)
                all_plot_x.append(final_y)
                all_plot_y.append(final_x)
            else:
                ax.plot(final_x, final_y, label=label, linestyle=style,
                        color=color, marker=marker, linewidth=linewidth,
                        solid_capstyle=solid_capstyle)
                all_plot_x.append(final_x)
                all_plot_y.append(final_y)

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

    # Final Styling
    if title:
        ax.set_title(_smart_label(title), fontsize=title_size)

    if swap_axes:
        if ylabel: ax.set_xlabel(_smart_label(ylabel), fontsize=label_size)
        if xlabel: ax.set_ylabel(_smart_label(xlabel), fontsize=label_size)
        if ylim:   ax.set_xlim(ylim)
        if xlim:   ax.set_ylim(xlim)
    else:
        if xlabel: ax.set_xlabel(_smart_label(xlabel), fontsize=label_size)
        if ylabel: ax.set_ylabel(_smart_label(ylabel), fontsize=label_size)
        if xlim:   ax.set_xlim(xlim)
        if ylim:   ax.set_ylim(ylim)

    ax.tick_params(axis='both', labelsize=tick_size)

    formatter = ticker.ScalarFormatter(useMathText=True)
    formatter.set_scientific(True)
    formatter.set_powerlimits((-2, 2))
    ax.xaxis.set_major_formatter(formatter)
    ax.yaxis.set_major_formatter(formatter)

    if scaling == 'constrained':
        ax.set_aspect('equal', adjustable='box')
    if labels:
        ax.legend(fontsize=legend_size)

    # Center Axes (move spines to origin)
    if center_axes:
        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')
        # Offset tick labels so they don't overlap the axes
        ax.tick_params(axis='x', pad=5)
        ax.tick_params(axis='y', pad=5)
        # Move axis labels to end of spines
        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', fontsize=label_size)
        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', fontsize=label_size)
    
    if grid:
        ax.grid(True)

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

    return ax