from __future__ import annotations

import os
import warnings
from collections import defaultdict
from itertools import chain
from typing import TYPE_CHECKING, Any

from matplotlib.pyplot import get_cmap
from monty.json import MSONable
from monty.serialization import loadfn
from palettable.colorbrewer.qualitative import Set1_9
from pymatgen.analysis.molecule_structure_comparator import CovalentRadius
from pymatgen.core import Element, Molecule, Site, Species
from pymatgen.util.string import unicodeify_species
from sklearn.preprocessing import LabelEncoder
from webcolors import html5_parse_legacy_color, html5_serialize_simple_color

from crystal_toolkit.settings import SETTINGS

if TYPE_CHECKING:
    from pymatgen.core.structure import SiteCollection

# element colors forked from pymatgen
module_dir = os.path.dirname(os.path.abspath(__file__))
EL_COLORS = loadfn(os.path.join(module_dir, "ElementColorSchemes.yaml"))


class Legend(MSONable):
    """Help generate a legend (colors and radii) for a Structure or Molecule such that colors and
    radii can be displayed for the appropriate species.

    Note that species themselves have a color (for example, Oxygen is typically red), but that we
    might also want to color-code by site properties (for example, magnetic moment), thus this class
    has to take into account both the species present and its context (the specific site the species
    is at) to correctly generate the legend.
    """

    default_color_scheme = SETTINGS.LEGEND_COLOR_SCHEME
    default_color = SETTINGS.LEGEND_FALLBACK_COLOR
    fallback_radius = SETTINGS.LEGEND_FALLBACK_RADIUS
    uniform_radius = SETTINGS.LEGEND_UNIFORM_RADIUS

    def __init__(
        self,
        site_collection: SiteCollection | Site,
        color_scheme: str = SETTINGS.LEGEND_COLOR_SCHEME,
        radius_scheme: str = SETTINGS.LEGEND_RADIUS_SCHEME,
        cmap: SETTINGS.LEGEND_CMAP = "coolwarm",
        cmap_range: tuple[float, float] | None = None,
    ) -> None:
        """Create a legend for a given SiteCollection to choose how to display colors and radii for
        the given sites and the species on those sites.

        If a site has a "display_color" or "display_radius" site
        property defined, this can be used to manually override the
        displayed colors and radii respectively.

        Args:
            site_collection: SiteCollection or, for convenience, a
            single site can be provided and this will be converted
            into a SiteCollection
            color_scheme: choose how to color-code species, one of
            "Jmol", "VESTA", "accessible" or a scalar site property
            (e.g. magnetic moment) or a categorical/string site
            property (e.g. Wyckoff label)
            radius_scheme: choose the radius for a species, one of
            "atomic", "specified_or_average_ionic", "covalent",
            "van_der_waals", "atomic_calculated", "uniform"
            cmap: only used if color_mode is set to a scalar site
            property, defines the matplotlib color map to use, by
            default is blue-white-red for negative to positive values
            cmap_range: only used if color_mode is set to a scalar site
            property, defines the minimum and maximum values of the
            color scape
        """
        if isinstance(site_collection, Site):
            site_collection = Molecule.from_sites([site_collection])

        site_prop_types = self.analyze_site_props(site_collection)

        self.allowed_color_schemes = [
            "VESTA",
            "Jmol",
            "accessible",
            *site_prop_types.get("scalar", []),
            *site_prop_types.get("categorical", []),
        ]

        self.allowed_radius_schemes = (
            "atomic",
            "specified_or_average_ionic",
            "covalent",
            "van_der_waals",
            "atomic_calculated",
            "uniform",
        )

        if color_scheme not in self.allowed_color_schemes:
            warnings.warn(
                f"Color scheme {color_scheme} not available, "
                f"falling back to {SETTINGS.LEGEND_COLOR_SCHEME}."
            )
            color_scheme = SETTINGS.LEGEND_COLOR_SCHEME

        # if color-coding by a scalar site property, determine minimum and
        # maximum values for color scheme, will default to be symmetric
        # about zero
        if color_scheme in site_prop_types.get("scalar", []) and not cmap_range:
            props = [
                p
                for p in site_collection.site_properties[color_scheme]
                if p is not None
            ]
            prop_max = max([abs(min(props)), max(props)])
            prop_min = -prop_max
            cmap_range = (prop_min, prop_max)

        el_colors = EL_COLORS.copy()
        el_colors.update(
            self.generate_accessible_color_scheme_on_the_fly(site_collection)
        )

        self.categorical_colors = self.generate_categorical_color_scheme_on_the_fly(
            site_collection, site_prop_types
        )

        self.el_colors = el_colors
        self.site_prop_types = site_prop_types
        self.site_collection = site_collection
        self.color_scheme = color_scheme
        self.radius_scheme = radius_scheme
        self.cmap = cmap
        self.cmap_range = cmap_range

    @staticmethod
    def generate_accessible_color_scheme_on_the_fly(
        site_collection: SiteCollection,
    ) -> dict[str, dict[str, tuple[int, int, int]]]:
        """E.g. for a color scheme more appropriate for people with color blindness.

        Args:
            site_collection: SiteCollection

        Returns: A dictionary in similar format to EL_COLORS
        """
        color_scheme = {}

        all_species = set(
            chain.from_iterable(comp for comp in site_collection.species_and_occu)
        )
        all_elements = sorted(sp.as_dict()["element"] for sp in all_species)

        # thanks to https://doi.org/10.1038/nmeth.1618
        palette = [
            (0, 0, 0),  # 0, black
            (230, 159, 0),  # 1, orange
            (86, 180, 233),  # 2, sky blue
            (0, 158, 115),  # 3, bluish green
            (240, 228, 66),  # 4, yellow
            (0, 114, 178),  # 5, blue
            (213, 94, 0),  # 6, vermilion
            (204, 121, 167),  # 7, reddish purple
            (255, 255, 255),  # 8, white
        ]

        # similar to CPK, mapping element to palette index
        preferred_colors = {
            "O": 6,
            "N": 2,
            "C": 0,
            "H": 8,
            "F": 3,
            "Cl": 3,
            "Fe": 1,
            "Br": 7,
            "I": 7,
            "P": 1,
            "S": 4,
        }

        if len(set(all_elements)) > len(palette):
            warnings.warn(
                "Too many distinct types of site to use an accessible color scheme, "
                "some sites will be given the default color."
            )

        preferred_elements_present = [
            el for el in all_elements if el in preferred_colors
        ]

        colors_assigned = []
        for el in preferred_elements_present:
            if preferred_colors[el] not in colors_assigned:
                color_scheme[el] = palette[preferred_colors[el]]
                colors_assigned.append(preferred_colors[el])

        remaining_elements = [el for el in all_elements if el not in color_scheme]
        remaining_palette = [
            c for idx, c in enumerate(palette) if idx not in colors_assigned
        ]

        for el in remaining_elements:
            if remaining_palette:
                color_scheme[el] = remaining_palette.pop()

        return {"accessible": color_scheme}

    @staticmethod
    def generate_categorical_color_scheme_on_the_fly(
        site_collection: SiteCollection, site_prop_types: dict[str, list[str]]
    ) -> dict[str, dict[str, tuple[int, int, int]]]:
        """E.g. for Wyckoff.

        Args:
            site_collection (SiteCollection): The sites to generate a color scheme for.
            site_prop_types (dict[str, list[str]]): The categorical site property types.

        Returns: A dictionary in similar format to EL_COLORS
        """
        color_scheme = {}

        palette = Set1_9.colors

        for site_prop_name in site_prop_types.get("categorical", []):
            props = site_collection.site_properties[site_prop_name]
            props[props is None] = "None"

            label_enc = LabelEncoder()
            label_enc.fit(props)
            transformed_props = label_enc.transform(props)

            # if we have more categories than available colors,
            # arbitrarily group some categories together
            if len(set(props)) > len(palette):
                warnings.warn(
                    "Too many categories for a complete categorical color scheme."
                )
            transformed_props = [
                p if p < len(palette) else -1 for p in transformed_props
            ]

            colors = {name: palette[p] for name, p in zip(props, transformed_props)}

            color_scheme[site_prop_name] = colors

        return color_scheme

    def get_color(self, sp: Species | Element, site: Site | None = None) -> str:
        """Get a color to render a specific species. Optionally, you can provide a site for context,
        since ...

        Args:
            sp: Species or Element
            site: Site

        Returns: Color
        """
        # allow manual override by user
        if (
            site
            and "display_color" in site.properties
            and site.properties["display_color"] is not None
        ):
            color = site.properties["display_color"]
            # TODO: next two lines due to change in API, will be removed
            if isinstance(color, list) and isinstance(color[0], str):
                color = color[0]
            if isinstance(color, list):
                return html5_serialize_simple_color(color)
            return html5_serialize_simple_color(html5_parse_legacy_color(color))

        if self.color_scheme in ("VESTA", "Jmol", "accessible"):
            el = sp.as_dict()["element"]
            color = self.el_colors[self.color_scheme].get(
                el, self.el_colors["Extras"].get(el, self.default_color)
            )

        elif self.color_scheme in self.site_prop_types.get("scalar", []):
            if not site:
                raise ValueError(
                    "Requires a site for context to get the appropriate site property."
                )

            prop = site.properties[self.color_scheme]

            if prop:
                cmap = get_cmap(self.cmap)

                # normalize in [0, 1] range, as expected by cmap
                prop_min, prop_max = self.cmap_range
                prop_normed = (prop - prop_min) / (prop_max - prop_min)

                color = [int(c * 255) for c in cmap(prop_normed)[0:3]]

            else:
                # fallback if site prop is None
                color = self.default_color

        elif self.color_scheme in self.site_prop_types.get("categorical", []):
            if not site:
                raise ValueError(
                    "Requires a site for context to get the appropriate site property."
                )

            prop = site.properties[self.color_scheme]

            color = self.categorical_colors[self.color_scheme].get(
                prop, self.default_color
            )

        else:
            raise ValueError(
                f"Unknown color for {sp} and color scheme {self.color_scheme}."
            )

        return html5_serialize_simple_color(color)

    def get_radius(self, sp: Species | Element, site: Site | None = None) -> float:
        # allow manual override by user
        if site and "display_radius" in site.properties:
            return site.properties["display_radius"]

        if self.radius_scheme not in self.allowed_radius_schemes:
            raise ValueError(
                f"Unknown radius scheme {self.radius_scheme}, "
                f"choose from: {self.allowed_radius_schemes}."
            )

        radius = None
        if self.radius_scheme == "uniform":
            radius = self.uniform_radius
        elif self.radius_scheme == "atomic":
            radius = float(sp.atomic_radius)
        elif (
            self.radius_scheme == "specified_or_average_ionic"
            and isinstance(sp, Species)
            and sp.oxi_state
        ):
            radius = float(sp.ionic_radius)
        elif self.radius_scheme == "specified_or_average_ionic":
            radius = float(sp.average_ionic_radius)
        elif self.radius_scheme == "covalent":
            el = str(getattr(sp, "element", sp))
            radius = float(CovalentRadius.radius[el])
        elif self.radius_scheme == "van_der_waals":
            radius = float(sp.van_der_waals_radius)
        elif self.radius_scheme == "atomic_calculated":
            radius = float(sp.atomic_radius_calculated)

        if (not radius) or (not isinstance(radius, float)):
            warnings.warn(
                f"Radius unknown for {sp} and strategy {self.radius_scheme}, "
                "setting to 0.5."
            )
            radius = self.fallback_radius

        return radius

    @staticmethod
    def analyze_site_props(site_collection: SiteCollection) -> dict[str, list[str]]:
        """A dictionary with keys "scalar", "matrix", "vector", "categorical"
        and values of a list of site property names corresponding to each type.
        """
        # (implicitly assumes all site props for a given key are same type)
        site_prop_names = defaultdict(list)
        for name, props in site_collection.site_properties.items():
            if isinstance(props[0], (int, float)):
                site_prop_names["scalar"].append(name)
            elif isinstance(props[0], list) and len(props[0]) == 3:
                if isinstance(props[0][0], list) and len(props[0][0]) == 3:
                    site_prop_names["matrix"].append(name)
                else:
                    site_prop_names["vector"].append(name)
            elif isinstance(props[0], str):
                site_prop_names["categorical"].append(name)
        return dict(site_prop_names)

    @staticmethod
    def get_species_str(sp: Species | Element) -> str:
        """Args:
            sp: Species or Element.

        Returns: string representation
        """
        # TODO: add roman numerals for oxidation state for ease of readability
        # and then move this to pymatgen string utils ...
        return unicodeify_species(str(sp))

    def get_legend(self) -> dict[str, Any]:
        # decide what we want the labels to be
        if self.color_scheme in ("Jmol", "VESTA", "accessible"):
            label = lambda site, sp: self.get_species_str(sp)
        elif self.color_scheme in self.site_prop_types.get("scalar", {}):
            label = lambda site, sp: f"{site.properties[self.color_scheme]:.2f}"
        elif self.color_scheme in self.site_prop_types.get("categorical", {}):
            label = lambda site, sp: f"{site.properties[self.color_scheme]}"
        else:
            raise ValueError(f"Color scheme {self.color_scheme} not known.")

        legend = defaultdict(list)

        # first get all our colors for different species
        for site in self.site_collection:
            for sp in site.species:
                legend[self.get_color(sp, site)].append(label(site, sp))

        legend = {key: ", ".join(sorted(set(val))) for key, val in legend.items()}

        color_options = []
        for site_prop_type in ("scalar", "categorical"):
            if site_prop_type in self.site_prop_types:
                for prop in self.site_prop_types[site_prop_type]:
                    color_options.append(prop)  # noqa: PERF402

        return {
            "composition": self.site_collection.composition.as_dict(),
            "colors": legend,
            "available_color_schemes": color_options,
        }
