"""
Color palettes for publiplots.

This module provides carefully curated color palettes optimized for
publication-ready visualizations, with a focus on pastel colors that
are visually appealing and print well.
"""

from typing import Dict, List, Optional, Union

from publiplots.themes.rcparams import resolve_param
import matplotlib.pyplot as plt
import numpy as np

from publiplots.utils import is_categorical

# =============================================================================
# Color Palettes
# =============================================================================

# Categorical Palettes
# =============================================================================

PASTEL_CATEGORICAL: List[str] = [
    "#75b375",  # Soft green
    "#8e8ec1",  # Soft purple
    "#eeaa58",  # Soft orange
    "#e67e7e",  # Soft red
    "#7ec5d9",  # Soft blue
    "#f0b0c4",  # Soft pink
    "#b8b88a",  # Soft olive
    "#c9a3cf",  # Soft lavender
    "#f4c896",  # Soft peach
    "#8fc9a8",  # Soft teal
    "#dba3a3",  # Soft rose
    "#9eb3d4",  # Soft periwinkle
]
"""
Pastel categorical palette with 12 distinct colors.

Use case: Categorical data with multiple distinct groups.
Works well for bar plots, scatter plots with categorical hue, and legends.
Colors are designed to be distinguishable while maintaining a soft, cohesive look.

Example:
    >>> import publiplots as pp
    >>> colors = pp.get_palette('pastel_categorical', n_colors=5)
    >>> pp.scatterplot(data=df, x='x', y='y', hue='category', palette=colors)
"""

PASTEL_CATEGORICAL_MINIMAL: List[str] = [
    "#75b375",  # Soft green
    "#8e8ec1",  # Soft purple
    "#eeaa58",  # Soft orange
    "#7ec5d9",  # Soft blue
]
"""
Minimal pastel categorical palette with 4 colors.

Use case: Small number of categories (2-4).
Optimized for clarity when fewer distinct groups need to be shown.

Example:
    >>> colors = pp.get_palette('pastel_categorical_minimal', n_colors=3)
"""

# =============================================================================
# Sequential Palettes
# =============================================================================

PASTEL_SEQUENTIAL_GREEN: List[str] = [
    "#e8f5e8",
    "#c8e6c8",
    "#a8d7a8",
    "#88c888",
    "#68b968",
    "#4aa44a",
]
"""
Sequential green palette (light to dark).

Use case: Continuous data ranging from low to high values.
Ideal for heatmaps, choropleth maps, and representing intensity.
"""

PASTEL_SEQUENTIAL_BLUE: List[str] = [
    "#e8f4f8",
    "#c8e4f0",
    "#a8d4e8",
    "#88c4e0",
    "#68b4d8",
    "#4aa4d0",
]
"""
Sequential blue palette (light to dark).

Use case: Continuous data, alternative to green for variety or when
representing different data types in the same figure.
"""

PASTEL_SEQUENTIAL_PURPLE: List[str] = [
    "#f0e8f5",
    "#e0c8e6",
    "#d0a8d7",
    "#c088c8",
    "#b068b9",
    "#a048a4",
]
"""
Sequential purple palette (light to dark).

Use case: Continuous data, provides good contrast with green/blue palettes
when multiple sequential scales are needed in complex visualizations.
"""

# =============================================================================
# Diverging Palettes
# =============================================================================

PASTEL_DIVERGING_RED_BLUE: List[str] = [
    "#e67e7e",  # Soft red
    "#f4a4a4",
    "#ffd4d4",
    "#f5f5f5",  # Neutral gray
    "#d4e4ff",
    "#a4c4f4",
    "#7ec5d9",  # Soft blue
]
"""
Diverging red-blue palette.

Use case: Data with a meaningful center point (e.g., zero, mean, or neutral value).
Perfect for log fold changes, correlations, or any data where deviation from
center in either direction has different meanings.

Example:
    >>> colors = pp.get_palette('pastel_diverging_red_blue', n_colors=7)
    >>> pp.heatmap(data=df, palette=colors, center=0)
"""

PASTEL_DIVERGING_GREEN_PURPLE: List[str] = [
    "#75b375",  # Soft green
    "#98c898",
    "#c8e6c8",
    "#f5f5f5",  # Neutral gray
    "#e8d8f0",
    "#c8a8d7",
    "#8e8ec1",  # Soft purple
]
"""
Diverging green-purple palette.

Use case: Alternative to red-blue for diverging data, especially useful
when red-green colorblindness is a concern or when red-blue is already
used for another data dimension.
"""

# =============================================================================
# Special Purpose Palettes
# =============================================================================

PASTEL_SIGNIFICANCE: Dict[str, str] = {
    "significant": "#75b375",      # Green for significant
    "not_significant": "#999999",  # Gray for not significant
    "highly_significant": "#4aa44a",  # Darker green for highly significant
}
"""
Significance palette for p-value visualization.

Use case: Highlighting statistically significant results in enrichment
analyses, differential expression, or other hypothesis testing visualizations.

Example:
    >>> colors = pp.PASTEL_SIGNIFICANCE
    >>> pp.barplot_enrichment(categories=terms, values=scores,
    ...                       pvalues=pvals, significance_colors=colors)
"""

PASTEL_POSITIVE_NEGATIVE: Dict[str, str] = {
    "positive": "#75b375",  # Green
    "negative": "#e67e7e",  # Red
    "neutral": "#eeaa58",   # Orange
}
"""
Positive/negative/neutral palette.

Use case: Data with clear positive and negative values (e.g., correlation,
fold change, effect size) where neutral represents ambiguous or intermediate state.
"""

# =============================================================================
# Palette Dictionary
# =============================================================================

PALETTES: Dict[str, Union[List[str], Dict[str, str]]] = {
    "pastel_categorical": PASTEL_CATEGORICAL,
    "pastel_categorical_minimal": PASTEL_CATEGORICAL_MINIMAL,
    "pastel_sequential_green": PASTEL_SEQUENTIAL_GREEN,
    "pastel_sequential_blue": PASTEL_SEQUENTIAL_BLUE,
    "pastel_sequential_purple": PASTEL_SEQUENTIAL_PURPLE,
    "pastel_diverging_red_blue": PASTEL_DIVERGING_RED_BLUE,
    "pastel_diverging_green_purple": PASTEL_DIVERGING_GREEN_PURPLE,
    "pastel_significance": PASTEL_SIGNIFICANCE,
    "pastel_positive_negative": PASTEL_POSITIVE_NEGATIVE,
}
"""Dictionary of all available palettes."""


# =============================================================================
# Functions
# =============================================================================

def get_palette(
    name: str,
    n_colors: Optional[int] = None,
    reverse: bool = False
) -> Union[List[str], Dict[str, str]]:
    """
    Get a color palette by name.

    Parameters
    ----------
    name : str
        Name of the palette. Available palettes:
        - 'pastel_categorical': 12 distinct pastel colors
        - 'pastel_categorical_minimal': 4 colors for minimal categories
        - 'pastel_sequential_green': Sequential green palette
        - 'pastel_sequential_blue': Sequential blue palette
        - 'pastel_sequential_purple': Sequential purple palette
        - 'pastel_diverging_red_blue': Red-blue diverging palette
        - 'pastel_diverging_green_purple': Green-purple diverging palette
        - 'pastel_significance': Significance levels (dict)
        - 'pastel_positive_negative': Positive/negative/neutral (dict)
    n_colors : int, optional
        Number of colors to return. If None, returns all colors in the palette.
        For sequential/diverging palettes, interpolates to get exact number.
        Not applicable for dictionary palettes.
    reverse : bool, default=False
        Whether to reverse the palette order.

    Returns
    -------
    Union[List[str], Dict[str, str]]
        List of color hex codes or dictionary mapping names to colors.

    Raises
    ------
    ValueError
        If palette name is not found or n_colors is invalid.

    Examples
    --------
    Get a categorical palette:
    >>> colors = get_palette('pastel_categorical', n_colors=5)
    >>> len(colors)
    5

    Get a sequential palette with interpolation:
    >>> colors = get_palette('pastel_sequential_green', n_colors=10)

    Get a significance palette (returns dict):
    >>> colors = get_palette('pastel_significance')
    >>> colors['significant']
    '#75b375'
    """
    if name not in PALETTES:
        available = ', '.join(PALETTES.keys())
        raise ValueError(f"Unknown palette '{name}'. Available palettes: {available}")

    palette = PALETTES[name]

    # Handle dictionary palettes
    if isinstance(palette, dict):
        if n_colors is not None:
            raise ValueError(
                f"Palette '{name}' is a dictionary palette. "
                "n_colors parameter is not applicable."
            )
        return palette

    # Handle list palettes
    if n_colors is None:
        result = list(palette)
    elif n_colors <= len(palette):
        # Return subset of palette
        result = list(palette[:n_colors])
    else:
        # Interpolate to get more colors
        result = _interpolate_colors(palette, n_colors)

    if reverse:
        result = result[::-1]

    return result


def _interpolate_colors(colors: List[str], n: int) -> List[str]:
    """
    Interpolate between colors to create a palette with n colors.

    Parameters
    ----------
    colors : List[str]
        List of hex color codes to interpolate between.
    n : int
        Number of colors to generate.

    Returns
    -------
    List[str]
        List of n interpolated hex color codes.
    """
    from matplotlib.colors import LinearSegmentedColormap, to_hex

    # Create a colormap from the colors
    cmap = LinearSegmentedColormap.from_list("custom", colors, N=n)

    # Sample n colors from the colormap
    indices = np.linspace(0, 1, n)
    interpolated = [to_hex(cmap(idx)) for idx in indices]

    return interpolated


def resolve_palette(
    palette: Optional[Union[str, List[str], Dict[str, str]]] = None,
    n_colors: Optional[int] = None
) -> Union[List[str], Dict[str, str]]:
    """
    Resolve a palette to actual colors, supporting both publiplots and seaborn palettes.

    This helper function translates palette specifications to actual color values.
    It checks if a string palette name is a publiplots palette first, then falls
    back to seaborn palettes if not found.

    Parameters
    ----------
    palette : str, list, dict, or None
        Palette specification:
        - None: Returns default palette (or default color for n_colors=1)
        - str: Palette name (checks publiplots first, then seaborn)
        - list: List of color hex codes (returned as-is)
        - dict: Mapping from categories to colors (returned as-is)
    n_colors : int, optional
        Number of colors to return. Only applicable for string palette names.

    Returns
    -------
    Union[List[str], Dict[str, str]]
        Resolved palette as list of colors or dictionary.

    Examples
    --------
    Get default color for single-color plots:
    >>> colors = resolve_palette(n_colors=1)
    >>> colors
    ['#5d83c3']

    Get default palette for multi-color plots:
    >>> colors = resolve_palette(n_colors=3)
    >>> len(colors)
    3

    Resolve publiplots palette:
    >>> colors = resolve_palette('pastel_categorical', n_colors=5)

    Resolve seaborn palette:
    >>> colors = resolve_palette('viridis', n_colors=5)

    Pass through color list:
    >>> colors = resolve_palette(['#ff0000', '#00ff00', '#0000ff'])
    ['#ff0000', '#00ff00', '#0000ff']
    """
    import seaborn as sns

    # Handle None: use default palette for multiple colors, default color for single
    if palette is None:
        default_palette = resolve_param("palette", None)
        default_color = resolve_param("color", None)

        if n_colors is not None and n_colors > 1:
            # Use default palette for multiple colors
            return get_palette(default_palette, n_colors=n_colors)
        else:
            # Use default color for single color
            return [default_color]

    # Handle list/dict: return as-is
    if isinstance(palette, (list, dict)):
        return palette

    # Handle string: check publiplots first, then seaborn
    if isinstance(palette, str):
        # Try publiplots palette first
        if palette in PALETTES:
            return get_palette(palette, n_colors=n_colors)

        # Fall back to seaborn palette
        try:
                return sns.color_palette(palette, n_colors=n_colors).as_hex()
        except Exception as e:
            raise ValueError(
                f"Unknown palette '{palette}'. Not found in publiplots palettes "
                f"({', '.join(PALETTES.keys())}) or seaborn palettes. Error: {e}"
            )

def resolve_palette_mapping(
    values: Optional[List[str]] = None,
    palette: Optional[Union[str, Dict, List]] = None,
) -> Dict[str, str]:
    """
    Resolve a palette mapping to actual colors.

    Parameters
    ----------
    values : List[str], optional
        List of values to map to colors.
    palette : str, dict, or list, optional
        Palette specification:
        - None: Returns default palette (or default color for n_colors=1)
        - str: Palette name (checks publiplots first, then seaborn)
        - list: List of color hex codes (returned as-is)
        - dict: Mapping from categories to colors (returned as-is)
    """
    if values is None:
        return {}
    if isinstance(palette, dict):
        return palette
    if not is_categorical(values):
        return palette # continuous mapping
    palette = resolve_palette(palette, n_colors=len(values))
    return {value: palette[i % len(palette)] for i, value in enumerate(values)}


def list_palettes() -> List[str]:
    """
    List all available palette names.

    Returns
    -------
    List[str]
        List of available palette names.

    Examples
    --------
    >>> palettes = list_palettes()
    >>> print(palettes)
    ['pastel_categorical', 'pastel_sequential_green', ...]
    """
    return list(PALETTES.keys())


def show_palette(name: str, n_colors: Optional[int] = None) -> None:
    """
    Display a color palette visually.

    Creates a simple visualization showing the colors in the palette.
    Useful for exploring and comparing palettes.

    Parameters
    ----------
    name : str
        Name of the palette to display.
    n_colors : int, optional
        Number of colors to display. If None, shows all colors.

    Examples
    --------
    >>> show_palette('pastel_categorical')
    >>> show_palette('pastel_sequential_green', n_colors=10)
    """
    palette = get_palette(name, n_colors)

    if isinstance(palette, dict):
        # Display dictionary palette
        colors_list = list(palette.values())
        labels = list(palette.keys())
        n = len(colors_list)
    else:
        colors_list = palette
        labels = [f"{i+1}" for i in range(len(colors_list))]
        n = len(colors_list)

    # Create visualization
    fig, ax = plt.subplots(figsize=(n * 0.8, 2))

    for i, (color, label) in enumerate(zip(colors_list, labels)):
        ax.add_patch(plt.Rectangle((i, 0), 1, 1, facecolor=color, edgecolor='black'))
        ax.text(i + 0.5, 0.5, label, ha='center', va='center',
                fontsize=10, rotation=90 if len(label) > 3 else 0)

    ax.set_xlim(0, n)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(f"Palette: {name}", fontsize=14, pad=10)

    plt.tight_layout()
    plt.show()
