"""Color related utilities for dashboard plots.

Color palettes from https://python-graph-gallery.com/color-palette-finder/

Material palettes:
- amber_material
- blue_grey_material
- blue_material
- brown_material
- cyan_material
- deep_orange_material
- deep_purple_material
- green_material
- grey_material
- indigo_material
- light_blue_material
- light_green_material
- lime_material
- orange_material
- pink_material
- purple_material
- red_material
- teal_material
- yellow_material
"""

import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
from pypalettes import load_cmap


def hex_to_rgb(hex_color: str) -> list[int]:
    """Convert hex color to RGB list.

    Args:
        hex_color: Hex color string (e.g., '#FF5733').

    Returns:
        List of [R, G, B] values (0-255).

    """
    hex_color = hex_color.lstrip("#")
    return [int(hex_color[i : i + 2], 16) for i in (0, 2, 4)]


def get_cmap(variable: str) -> mcolors.Colormap:
    """Get a color palette by a "data" variable.

    Each variable (meaning of the data) should be associated with another color palette when plotting.
    This function should help to standardize the color palettes used for each variable type.

    The variable can be any string, descriptive names are recommended.

    Args:
        variable: The variable to load a palette for.

    Returns:
        A list of hex color strings.

    """
    material_palettes = [
        "amber_material",
        "blue_grey_material",
        "blue_material",
        "brown_material",
        "cyan_material",
        "deep_orange_material",
        "deep_purple_material",
        "green_material",
        "grey_material",
        "indigo_material",
        "light_blue_material",
        "light_green_material",
        "lime_material",
        "orange_material",
        "pink_material",
        "purple_material",
        "red_material",
        "teal_material",
        "yellow_material",
    ]
    # Fuzzy map from variable type to palette name
    material_idx = sum(ord(c) for c in variable) % len(material_palettes)
    palette_name = material_palettes[material_idx]
    cmap = load_cmap(name=palette_name)
    return cmap


def get_palette(variable: str, n_colors: int) -> list[str]:
    """Get a color palette by a "data" variable.

    Each variable (meaning of the data) should be associated with another color palette when plotting.
    This function should help to standardize the color palettes used for each variable type.

    The variable can be any string, descriptive names are recommended.

    Args:
        variable: The variable to load a palette for.
        n_colors: The number of colors to return.

    Returns:
        A list of hex color strings.

    """
    # Hardcode some common variables to specific colormaps
    if variable == "ERA5":
        cmap = load_cmap(name="blue_material").resampled(n_colors)
    elif variable == "ArcticDEM":
        cmap = load_cmap(name="deep_purple_material").resampled(n_colors)
    elif variable == "AlphaEarth":
        cmap = load_cmap(name="green_material").resampled(n_colors)
    else:
        cmap = get_cmap(variable).resampled(n_colors)
    colors = [mcolors.to_hex(cmap(i)) for i in range(cmap.N)]
    return colors


def generate_unified_colormap(
    settings: dict,
) -> tuple[mcolors.ListedColormap, mcolors.ListedColormap, list[str]]:
    """Generate unified colormaps for all plotting libraries.

    This function creates consistent color schemes across Matplotlib/Ultraplot,
    Folium/Leaflet, and Altair/Vega-Lite by determining the task type and number
    of classes from the settings, then generating appropriate colormaps for each library.

    Args:
        settings: Settings dictionary containing task type, classes, and other configuration.

    Returns:
        Tuple of (matplotlib_cmap, folium_cmap, altair_colors) where:
        - matplotlib_cmap: matplotlib ListedColormap object
        - folium_cmap: matplotlib ListedColormap object (for geopandas.explore)
        - altair_colors: list of hex color strings for Altair

    """
    # Determine task type and number of classes from settings
    task = settings.get("task", "binary")
    n_classes = len(settings.get("classes", []))

    # Check theme
    is_dark_theme = st.context.theme.type == "dark"

    # Define base colormaps for different tasks
    if task == "binary":
        # For binary: use a simple two-color scheme
        if is_dark_theme:
            base_colors = ["#1f77b4", "#ff7f0e"]  # Blue and orange for dark theme
        else:
            base_colors = [
                "#3498db",
                "#e74c3c",
            ]  # Brighter blue and red for light theme
    else:
        # For multi-class: use a sequential colormap
        # Use matplotlib's viridis colormap
        cmap = plt.get_cmap("viridis")
        # Sample colors evenly across the colormap
        indices = np.linspace(0.1, 0.9, n_classes)  # Avoid extreme ends
        base_colors = [mcolors.rgb2hex(cmap(idx).tolist()[:3]) for idx in indices]

    # Create matplotlib colormap (for ultraplot and geopandas)
    matplotlib_cmap = mcolors.ListedColormap(base_colors)

    # Create Folium/Leaflet colormap (geopandas.explore uses matplotlib colormaps)
    folium_cmap = mcolors.ListedColormap(base_colors)

    # Create Altair color list (Altair uses hex color strings in range)
    altair_colors = base_colors

    return matplotlib_cmap, folium_cmap, altair_colors
