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,
            # --- Plot Labels ---
            title=None, xlabel=None, ylabel=None, labels=None,
            # --- Styling ---
            colors=None, line_styles=None, linewidth=1, scaling=None,
            solid_capstyle='butt',
            # --- Typography ---
            fontsize=14, title_size=None, label_size=None,
            tick_size=None, legend_size=None,
            # --- Axis Options ---
            swap_axes=False,
            # --- Limits ---
            xlim=None, ylim=None,
            # --- Resolution ---
            resolution=400,
            # --- Display settings ---
            show=True,width=800, height=600,
            # --- Break Control ---
            _break_=None):

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

    Parameters:
        exprs        : Expression, tuple, or list of expressions to plot
                        1. (x,y) for parametric plot
                        2. [x,y] for standard   plot
        var          : Symbol or tuple defining the variable and range
        labels       : Labels for each curve (default=None)
        line_styles  : Line styles or markers for each curve (default=None)
        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)
        scaling      : 'constrained'→ equal axis units(default=None)
        solid_capstyle : 'butt', 'round', or 'projecting' (default='butt')
                         - 'butt': Ends line exactly at data point (Good for connecting beams).
                         - 'round': Rounded ends.
                         - 'projecting': Square ends that extend slightly beyond data point.

    Returns:
        matplotlib Axes
            The generated 2D plot 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)}$"

    # Handle Parametric Logic
    def _get_parametric_data(expr, x_sym, x_vals_sample):
        """
        Handles:
          1. (x_expr, y_expr, (t, min, max)) -> Custom range
          2. (x_expr, y_expr) -> Uses global range
        """
        # Case A: Custom Range (x, y, (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)
            
            # Broadcast scalars
            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()

        # Case B: Global Range (x, y)
        elif len(expr) == 2:
            x_data, y_data = expr
            
            # Check if symbolic
            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)

            # Check for breaks
            if expr in _break_:
                y_data = _break_kinks(y_data)
                
            return x_data, y_data
        
        return None, None

    # 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.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.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_ = []
    
    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 extraction
        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
        if isinstance(expr, Matrix):
            expr = np.array(expr).astype(np.float64).flatten()

        # DISPATCH LOGIC
        final_x = None
        final_y = None

        # A. Parametric (Tuple/List)
        if isinstance(expr, (tuple, list)):
             final_x, final_y = _get_parametric_data(expr, x_sym, x_vals_sample)
        
        # B. Standard (Expression)
        else:
             final_x, final_y = _get_standard_data(expr, x_sym, x_vals_sample)

        # C. Handle Markers vs Lines
        if style in marker_symbols:
            marker = style
            style = ''

        # D. Plot & Collect
        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)
        # Apply Manual Limits (Override Padding)
        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)
        # Apply Manual Limits (Override Padding)
        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)
    ax.grid(True)

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

    return ax