"""
Visualization module for MedTDA.

This module provides plotting functions for persistence diagrams, barcodes,
Betti curves, and persistence landscapes.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from typing import Optional, Union, Dict, List, Tuple
import warnings

# Import vectorization functions for landscape computation
from . import vectorizers

try:
    import seaborn as sns
    SEABORN_AVAILABLE = True
except ImportError:
    SEABORN_AVAILABLE = False
    warnings.warn(
        "Seaborn not available. Install with 'pip install seaborn' for enhanced visualizations."
    )


# Set default theme on module import
if SEABORN_AVAILABLE:
    sns.set_theme(style='darkgrid')


def plot_persistence_diagram(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    markersize: int = 5,
    alpha: float = 1,
    show_diagonal: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid'
) -> plt.Axes:
    """
    Plot persistence diagram.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        If array: (n_features, 2) array of (birth, death) pairs.
        If dict: Dictionary with keys like 'H0', 'H1', etc., mapping to barcode arrays.
    dimensions : list of int, optional
        Which dimensions to plot. If None, plots all available dimensions.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on. If None, creates new figure.
    color_palette : list, optional
        Color palette for different dimensions. If None, uses viridis palette.
    markersize : int, default=5
        Size of markers.
    alpha : float, default=0.6
        Transparency of markers (0-1).
    show_diagonal : bool, default=True
        Whether to show the diagonal line (birth = death).
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title. If None, uses default title.
    theme : str, default='darkgrid'
        Seaborn theme style ('darkgrid', 'whitegrid', 'dark', 'white', 'ticks').
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda import BarcodeExtractor
    >>> from medtda.plotting import plot_persistence_diagram
    >>> extractor = BarcodeExtractor()
    >>> barcodes = extractor.execute(image)
    >>> plot_persistence_diagram(barcodes)
    >>> plt.show()
    """
    # Set seaborn theme first (before any other operations)
    if SEABORN_AVAILABLE:
        sns.set_theme(style=theme)
    
    # Set default color palette if not provided
    if color_palette is None:
        if SEABORN_AVAILABLE:
            color_palette = sns.color_palette("viridis")
        else:
            color_palette = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
    
    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 6))
    
    # Convert single barcode to dict format
    if isinstance(barcodes, np.ndarray):
        barcodes = {'H0': barcodes}
    
    # Determine which dimensions to plot
    if dimensions is None:
        # Extract dimension numbers from keys
        available_dims = []
        for key in barcodes.keys():
            if key.startswith('H'):
                dim = int(key[1:])
                available_dims.append(dim)
        dimensions = sorted(available_dims)
    
    # Prepare list of diagrams (filtering finite values)
    diagrams = []
    for dim in dimensions:
        key = f'H{dim}'
        if key in barcodes and barcodes[key].shape[0] > 0:
            barcode = barcodes[key]
            # Filter to finite values only
            finite_mask = np.isfinite(barcode[:, 1])
            if np.any(finite_mask):
                diagrams.append(barcode[finite_mask])
            else:
                diagrams.append(np.empty((0, 2)))
        else:
            diagrams.append(np.empty((0, 2)))
    
    # Plot persistence diagrams
    for dim, diagram in enumerate(diagrams):
        if diagram.size == 0:
            continue
        ax.scatter(diagram[:, 0], diagram[:, 1], color=color_palette[dim], 
                  label=f'$H_{{{dim}}}$', alpha=alpha, s=markersize)
    
    # Add diagonal line
    if show_diagonal:
        non_empty = [d for d in diagrams if d.size > 0]
        if non_empty:
            min_birth: float = float(min(np.min(d[:, 0]) for d in non_empty))
            max_val = max(np.max(d) for d in non_empty) * 1.1
            ax.plot([min_birth, max_val], [min_birth, max_val], 'k--', lw=1)
    
    # Labels and styling
    ax.set_xlabel('Birth')
    ax.set_ylabel('Death')
    ax.set_title(title if title else 'Persistence Diagram')
    ax.grid(True)
    
    if legend:
        ax.legend(loc='lower right')
    
    return ax


def plot_barcode(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    linewidth: float = 0.5,
    max_bars: Optional[int] = None,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid'
) -> plt.Axes:
    """
    Plot persistence barcode.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        If array: (n_features, 2) array of (birth, death) pairs.
        If dict: Dictionary with keys like 'H0', 'H1', etc., mapping to barcode arrays.
    dimensions : list of int, optional
        Which dimensions to plot. If None, plots all available dimensions.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on. If None, creates new figure.
    color_palette : list, optional
        Color palette for different dimensions. If None, uses viridis palette.
    linewidth : float, default=0.5
        Width of barcode lines.
    max_bars : int, optional
        Maximum number of bars to show per dimension (shows longest-lived features).
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda.plotting import plot_barcode
    >>> plot_barcode(barcodes, max_bars=50)
    >>> plt.show()
    """
    # Set seaborn theme first (before any other operations)
    if SEABORN_AVAILABLE:
        sns.set_theme(style=theme)
    
    # Set default color palette if not provided
    if color_palette is None:
        if SEABORN_AVAILABLE:
            color_palette = sns.color_palette("viridis")
        else:
            color_palette = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
    
    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))
    
    # Convert single barcode to dict format
    if isinstance(barcodes, np.ndarray):
        barcodes = {'H0': barcodes}
    
    # Determine which dimensions to plot
    if dimensions is None:
        available_dims = []
        for key in barcodes.keys():
            if key.startswith('H'):
                dim = int(key[1:])
                available_dims.append(dim)
        dimensions = sorted(available_dims)
    
    # Define gap between different homology dimensions
    gap_between_groups = 10
    
    # Track current y position to avoid overlap
    current_y = 0
    
    # Plot each dimension
    for dim in dimensions:
        key = f'H{dim}'
        if key not in barcodes:
            continue
        
        barcode = barcodes[key]
        if barcode.shape[0] == 0:
            continue
        
        # Limit number of bars if specified
        if max_bars and len(barcode) > max_bars:
            # Keep longest-lived features
            persistence = barcode[:, 1] - barcode[:, 0]
            persistence[~np.isfinite(persistence)] = 0  # Handle infinite deaths
            idx = np.argsort(persistence)[-max_bars:]
            barcode = barcode[idx]
        
        # Vectorized plotting - much faster than looping
        n_bars = len(barcode)
        y_positions = np.arange(n_bars) + current_y
        births = barcode[:, 0]
        deaths = np.where(np.isfinite(barcode[:, 1]), barcode[:, 1], np.maximum(births + 1, 5))
        
        # Plot all bars for this dimension at once
        ax.hlines(y_positions, births, deaths, colors=color_palette[dim], lw=linewidth,
                 label=f'$H_{{{dim}}}$')
        
        # Update y position for next dimension (add number of bars + gap)
        current_y += n_bars + gap_between_groups
    
    # Labels and styling
    ax.set_xlabel('Filtration Value')
    ax.set_ylabel('Barcodes')
    ax.set_title(title if title else 'Persistent Barcodes')
    ax.grid(True)
    
    if legend:
        ax.legend(loc='upper left', bbox_to_anchor=(0, 0.95))
    
    return ax


def _plot_vectorization_curve(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    vectorization_func,
    vectorization_params: dict,
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    fill: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid',
    xlabel: str = 'Sample Points',
    ylabel: str = 'Value'
) -> plt.Axes:
    """
    Internal helper function for plotting vectorization curves.
    
    This function handles the common logic for plotting curve-based vectorizations
    like Betti curves, entropy summaries, lifespan curves, and silhouettes.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        Barcodes to plot.
    vectorization_func : callable
        Vectorization function from vectorizers module.
    vectorization_params : dict
        Parameters to pass to the vectorization function.
    dimensions : list of int, optional
        Which dimensions to plot.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on.
    color_palette : list, optional
        Color palette for different dimensions.
    fill : bool, default=True
        Whether to fill area under curve.
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
    xlabel : str, default='Sample Points'
        X-axis label.
    ylabel : str, default='Value'
        Y-axis label.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
    """
    # Set seaborn theme first (before any other operations)
    if SEABORN_AVAILABLE:
        sns.set_theme(style=theme)
    
    # Set default color palette if not provided
    if color_palette is None:
        if SEABORN_AVAILABLE:
            color_palette = sns.color_palette("viridis")
        else:
            color_palette = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
    
    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))
    
    # Convert single barcode to dict format
    if isinstance(barcodes, np.ndarray):
        barcodes = {'H0': barcodes}
    
    # Determine which dimensions to plot
    if dimensions is None:
        available_dims = []
        for key in barcodes.keys():
            if key.startswith('H'):
                dim = int(key[1:])
                available_dims.append(dim)
        dimensions = sorted(available_dims)
    
    # Get resolution from params
    resolution = vectorization_params.get('resolution', 100)
    x_values = np.arange(resolution)
    
    # Plot curve for each dimension
    for dim in dimensions:
        key = f'H{dim}'
        if key not in barcodes or barcodes[key].shape[0] == 0:
            continue
        
        barcode = barcodes[key]
        
        # Compute vectorization
        curve_values = vectorization_func(barcode, **vectorization_params)
        
        # Plot and fill area under curve
        ax.plot(x_values, curve_values, color=color_palette[dim], label=f'$dim={dim}$')
        if fill:
            ax.fill_between(x_values, curve_values, color=color_palette[dim], alpha=0.3)
    
    # Labels and styling
    ax.set_title(title if title is not None else '')
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.grid(True)
    
    if legend:
        ax.legend(loc='upper left')
    
    return ax


def plot_betti_curve(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    resolution: int = 100,
    fill: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid'
) -> plt.Axes:
    """
    Plot Betti curves.
    
    Betti curve shows the number of topological features (Betti numbers)
    as a function of filtration value.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        Barcodes to plot.
    dimensions : list of int, optional
        Which dimensions to plot.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on.
    color_palette : list, optional
        Color palette for different dimensions. If None, uses viridis palette.
    resolution : int, default=100
        Number of sample points for curve resolution.
    fill : bool, default=True
        Whether to fill area under curve.
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda.plotting import plot_betti_curve
    >>> plot_betti_curve(barcodes)
    >>> plt.show()
    """
    return _plot_vectorization_curve(
        barcodes=barcodes,
        vectorization_func=vectorizers.betti_curve,
        vectorization_params={'resolution': resolution},
        dimensions=dimensions,
        ax=ax,
        color_palette=color_palette,
        fill=fill,
        legend=legend,
        title=title if title else 'Betti Curve',
        theme=theme,
        xlabel='Filtration Value',
        ylabel='Betti Number'
    )


def plot_entropy_summary(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    resolution: int = 100,
    fill: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid'
) -> plt.Axes:
    """
    Plot entropy summary curves.
    
    Entropy summary provides a statistical representation based on persistence entropy.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        Barcodes to plot.
    dimensions : list of int, optional
        Which dimensions to plot.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on.
    color_palette : list, optional
        Color palette for different dimensions.
    resolution : int, default=100
        Number of sample points.
    fill : bool, default=True
        Whether to fill area under curve.
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda.plotting import plot_entropy_summary
    >>> plot_entropy_summary(barcodes)
    >>> plt.show()
    """
    return _plot_vectorization_curve(
        barcodes=barcodes,
        vectorization_func=vectorizers.entropy_summary,
        vectorization_params={'resolution': resolution},
        dimensions=dimensions,
        ax=ax,
        color_palette=color_palette,
        fill=fill,
        legend=legend,
        title=title if title else 'Entropy Summary Function',
        theme=theme,
        xlabel='Sample Points',
        ylabel='Entropy'
    )


def plot_lifespan(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    resolution: int = 100,
    fill: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid'
) -> plt.Axes:
    """
    Plot persistence lifespan curves.
    
    Lifespan curve represents the normalized lifespan distribution of topological features.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        Barcodes to plot.
    dimensions : list of int, optional
        Which dimensions to plot.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on.
    color_palette : list, optional
        Color palette for different dimensions.
    resolution : int, default=100
        Number of sample points.
    fill : bool, default=True
        Whether to fill area under curve.
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda.plotting import plot_lifespan
    >>> plot_lifespan(barcodes)
    >>> plt.show()
    """
    return _plot_vectorization_curve(
        barcodes=barcodes,
        vectorization_func=vectorizers.persistence_lifespan,
        vectorization_params={'resolution': resolution},
        dimensions=dimensions,
        ax=ax,
        color_palette=color_palette,
        fill=fill,
        legend=legend,
        title=title if title else 'Persistence Lifespan Curve',
        theme=theme,
        xlabel='Sample Points',
        ylabel='Lifespan'
    )


def plot_silhouette(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    resolution: int = 100,
    weight: float = 1.0,
    fill: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid'
) -> plt.Axes:
    """
    Plot persistence silhouette curves.
    
    Silhouette provides a weighted summary of the persistence diagram.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        Barcodes to plot.
    dimensions : list of int, optional
        Which dimensions to plot.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on.
    color_palette : list, optional
        Color palette for different dimensions.
    resolution : int, default=100
        Number of sample points.
    weight : float, default=1.0
        Weight parameter for silhouette computation.
    fill : bool, default=True
        Whether to fill area under curve.
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda.plotting import plot_silhouette
    >>> plot_silhouette(barcodes)
    >>> plt.show()
    """
    return _plot_vectorization_curve(
        barcodes=barcodes,
        vectorization_func=vectorizers.persistence_silhouette,
        vectorization_params={'resolution': resolution, 'weight': weight},
        dimensions=dimensions,
        ax=ax,
        color_palette=color_palette,
        fill=fill,
        legend=legend,
        title=title if title else 'Persistence Silhouette',
        theme=theme,
        xlabel='Sample Points',
        ylabel='Silhouette Value'
    )


def plot_landscape(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimension: int = 1,
    num_landscapes: int = 5,
    resolution: int = 100,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    fill: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid'
) -> plt.Axes:
    """
    Plot persistence landscape.
    
    Persistence landscape is a functional representation of persistence barcodes
    obtained by "stacking" tent functions.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        Barcodes to plot.
    dimension : int, default=1
        Which homological dimension to plot.
    num_landscapes : int, default=5
        Number of landscape functions to plot (k=1, 2, ..., num_landscapes).
    resolution : int, default=100
        Number of points for curve resolution.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on.
    color_palette : list, optional
        Color palette for different landscape levels. If None, uses viridis palette.
    fill : bool, default=True
        Whether to fill area under landscape curves.
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda.plotting import plot_landscape
    >>> plot_landscape(barcodes, dimension=1, num_landscapes=5)
    >>> plt.show()
    """
    # Set seaborn theme first (before any other operations)
    if SEABORN_AVAILABLE:
        sns.set_theme(style=theme)
    
    # Set default color palette if not provided
    if color_palette is None:
        if SEABORN_AVAILABLE:
            color_palette = sns.color_palette("viridis", num_landscapes)
        else:
            # Fallback to matplotlib colormap
            try:
                colors_array = plt.colormaps['viridis'](np.linspace(0.2, 0.8, num_landscapes))
                color_palette = [tuple(c) for c in colors_array]
            except AttributeError:
                colors_array = plt.cm.get_cmap('viridis')(np.linspace(0.2, 0.8, num_landscapes))
                color_palette = [tuple(c) for c in colors_array]
    
    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))
    
    # Get barcode for specified dimension
    if isinstance(barcodes, dict):
        key = f'H{dimension}'
        if key not in barcodes:
            raise ValueError(f"Dimension {dimension} not found in barcodes.")
        barcode = barcodes[key]
    else:
        barcode = barcodes
    
    if barcode.shape[0] == 0:
        warnings.warn(f"No features in dimension {dimension}.")
        return ax
    
    # Remove infinite death times
    finite_mask = np.isfinite(barcode[:, 1])
    barcode = barcode[finite_mask]
    
    if len(barcode) == 0:
        warnings.warn(f"No finite features in dimension {dimension}.")
        return ax
    
    # Compute persistence landscape using vectorization
    landscape_vector = vectorizers.persistence_landscape(
        barcode, 
        resolution=resolution, 
        num_landscapes=num_landscapes
    )
    
    # Determine filtration range for x-axis
    min_val: float = float(np.min(barcode))
    max_val: float = float(np.max(barcode))
    x_values = np.linspace(min_val, max_val, resolution)
    
    # Plot each landscape with region overlays
    for i in range(num_landscapes):
        # Extract landscape values for this level
        landscape_k = landscape_vector[i * resolution:(i + 1) * resolution]
        
        # Fill area under landscape first (if enabled)
        if fill:
            ax.fill_between(x_values, 0, landscape_k, color=color_palette[i], alpha=0.4)
        
        # Plot landscape line on top
        ax.plot(x_values, landscape_k, color=color_palette[i], label=f'$L_{{{i+1}}}$')
    
    # Labels and styling
    ax.set_title(title if title else f'Persistent Landscapes $[dim={dimension}]$')
    ax.grid(True)
    
    # Use colorbar instead of legend when there are many landscapes
    if legend:
        if num_landscapes <= 5:
            # Use traditional legend for few landscapes
            ax.legend()
        else:
            # Use discrete colorbar for many landscapes
            from matplotlib.colors import ListedColormap, BoundaryNorm
            from matplotlib.cm import ScalarMappable
            
            # Create discrete colormap
            cmap = ListedColormap(color_palette)
            bounds = np.arange(num_landscapes + 1)
            norm = BoundaryNorm(bounds, cmap.N)
            
            # Create colorbar with only min and max labels
            sm = ScalarMappable(cmap=cmap, norm=norm)
            sm.set_array([])
            cbar = plt.colorbar(sm, ax=ax, ticks=[0.5, num_landscapes - 0.5], 
                              spacing='proportional', pad=0.02)
            cbar.set_label('Landscape Level', rotation=270, labelpad=15)
            cbar.ax.set_yticklabels([f'$L_1$', f'$L_{{{num_landscapes}}}$'])
    
    return ax


def plot_tropical_coordinates(
    barcodes: Union[np.ndarray, Dict[str, np.ndarray]],
    dimensions: Optional[List[int]] = None,
    ax: Optional[plt.Axes] = None,
    color_palette: Optional[List] = None,
    bar_width: float = 0.25,
    alpha: float = 0.8,
    log_scale: bool = True,
    legend: bool = True,
    title: Optional[str] = None,
    theme: str = 'darkgrid',
    r: float = 28
) -> plt.Axes:
    """
    Plot tropical coordinates as a bar chart.
    
    Tropical coordinates provide polynomial invariants of persistence diagrams
    using tropical algebra.
    
    Parameters
    ----------
    barcodes : numpy.ndarray or dict
        Barcodes to plot.
    dimensions : list of int, optional
        Which dimensions to plot.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on.
    color_palette : list, optional
        Color palette for different dimensions.
    bar_width : float, default=0.35
        Width of individual bars.
    alpha : float, default=0.8
        Transparency of bars (0-1).
    log_scale : bool, default=True
        Whether to use logarithmic scale for y-axis.
    legend : bool, default=True
        Whether to show legend.
    title : str, optional
        Plot title.
    theme : str, default='darkgrid'
        Seaborn theme style.
    r : float, default=28
        Parameter for tropical coordinate computation.
        
    Returns
    -------
    matplotlib.axes.Axes
        The axes object with the plot.
        
    Examples
    --------
    >>> from medtda.plotting import plot_tropical_coordinates
    >>> plot_tropical_coordinates(barcodes)
    >>> plt.show()
    """
    # Set seaborn theme first (before any other operations)
    if SEABORN_AVAILABLE:
        sns.set_theme(style=theme)
    
    # Set default color palette if not provided
    if color_palette is None:
        if SEABORN_AVAILABLE:
            color_palette = sns.color_palette("viridis")
        else:
            color_palette = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
    
    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))
    
    # Convert single barcode to dict format
    if isinstance(barcodes, np.ndarray):
        barcodes = {'H0': barcodes}
    
    # Determine which dimensions to plot
    if dimensions is None:
        available_dims = []
        for key in barcodes.keys():
            if key.startswith('H'):
                dim = int(key[1:])
                available_dims.append(dim)
        dimensions = sorted(available_dims)
    
    # Compute tropical coordinates for each dimension
    tropical_coords = {}
    for dim in dimensions:
        key = f'H{dim}'
        if key not in barcodes or barcodes[key].shape[0] == 0:
            continue
        
        barcode = barcodes[key]
        coords = vectorizers.persistence_tropical_coordinates(barcode, r=r)
        tropical_coords[dim] = coords
    
    if len(tropical_coords) == 0:
        warnings.warn("No valid barcodes to plot.")
        return ax
    
    # Number of coordinates (typically 7)
    num_coords = len(next(iter(tropical_coords.values())))
    x = np.arange(num_coords)
    
    # Calculate bar positions for grouped bars
    num_dims = len(tropical_coords)
    width_per_dim = bar_width
    total_width = width_per_dim * num_dims
    
    # Plot bars for each dimension
    for idx, dim in enumerate(sorted(tropical_coords.keys())):
        coords = tropical_coords[dim]
        
        # Calculate offset for this dimension
        offset = (idx - num_dims / 2) * width_per_dim + width_per_dim / 2
        bar_positions = x + offset
        
        # Plot bars
        ax.bar(bar_positions, coords, width=width_per_dim, 
               color=color_palette[dim], label=f'$dim={dim}$', alpha=alpha)
    
    # Apply logarithmic scale to y-axis if requested
    if log_scale:
        ax.set_yscale('log')
    
    # Labels and styling
    ax.set_title(title if title else 'Tropical Coordinates')
    ax.set_xlabel('Coordinate')
    ax.set_ylabel('Value')
    ax.set_xticks(x)
    ax.set_xticklabels([f'{i}' for i in range(num_coords)])
    ax.grid(True, alpha=0.3)
    
    if legend:
        ax.legend()
    
    return ax
