import os
import platform
from enum import Enum
from pathlib import Path
from typing import Any, Literal

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import rcParams
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from numpy.typing import ArrayLike, NDArray


def set_legend(
    ax: Axes | None = None,
    unique: bool = False,
    shift: tuple[int, int] = (0, 0),
    box: bool = False,
    **kwargs: Any,
) -> None:
    """Configure and add a legend to a matplotlib axes.

    This function provides convenient legend positioning including special
    "outside" locations, optional shifting, and unique label filtering.

    Parameters
    ----------
    ax : Axes or None, optional
        The axes to add the legend to. If None, uses current axes.
    unique : bool, optional
        If True, filters out duplicate labels keeping only the first occurrence.
        Default is False.
    shift : tuple of int, optional
        A (dx, dy) shift to apply to the legend's bbox_to_anchor position.
        Default is (0, 0).
    box : bool, optional
        If True, draws a fancy box around the legend with transparency.
        If False, removes the legend frame. Default is False.
    **kwargs : Any
        Additional keyword arguments passed to ax.legend(). Special loc values
        include "outside right", "outside upper"/"outside top", and
        "outside lower"/"outside bottom" for positioning legends outside the plot area.
    """
    if box:
        rcParams["legend.fancybox"] = True
        rcParams["legend.frameon"] = True
        rcParams["legend.framealpha"] = 0.8
    else:
        rcParams["legend.frameon"] = False
        rcParams["legend.fancybox"] = False
        rcParams["legend.framealpha"] = None
        rcParams["legend.edgecolor"] = "inherit"

    if ax is None:
        ax = plt.gca()

    handles, labels = ax.get_legend_handles_labels()

    loc = kwargs.get("loc", None)
    if loc == "outside right":
        kwargs["loc"] = "center left"
        kwargs["bbox_to_anchor"] = np.array([1.0, 0.5])
        kwargs["bbox_transform"] = ax.transAxes
    elif loc in ["outside upper", "outside top"]:
        kwargs["loc"] = "lower center"
        kwargs["bbox_to_anchor"] = np.array([0.5, 1.0])
        kwargs["bbox_transform"] = ax.transAxes
        if not kwargs.get("ncol"):
            kwargs["ncol"] = len(labels)
    elif loc in ["outside lower", "outside bottom"]:
        kwargs["loc"] = "lower center"
        kwargs["bbox_to_anchor"] = np.array([0.5, -0.3])
        kwargs["bbox_transform"] = ax.transAxes
        if not kwargs.get("ncol"):
            kwargs["ncol"] = len(labels)

    bba = kwargs.get("bbox_to_anchor")
    if bba is not None:
        new_bba = np.asarray(bba) + shift
        kwargs["bbox_to_anchor"] = new_bba

    if unique:
        leg_uniq = {k: v for k, v in zip(labels, handles)}
        ax.legend(leg_uniq.values(), leg_uniq.keys(), **kwargs)
    else:
        ax.legend(**kwargs)


def set_grid(ax: Axes | None = None, seaborn_style: bool = True) -> None:
    """Add a grid to a matplotlib axes with optional seaborn styling.

    Parameters
    ----------
    ax : Axes or None, optional
        The axes to add the grid to. If None, uses current axes.
    seaborn_style : bool, optional
        If True, applies a seaborn-like grid style with light grey background
        and white grid lines. If False, uses matplotlib's default grid style.
        Default is True.
    """
    if ax is None:
        ax = plt.gca()

    if seaborn_style:
        ax.patch.set_facecolor("lightgrey")
        ax.patch.set_alpha(0.5)
        ax.grid(visible=True, color="white")
        ax.set_axisbelow(True)
    else:
        ax.grid(visible=True)


def set_conditional_logscale(
    ax: Axes,
    data: ArrayLike,
    which_axis: Literal["x", "y"] = "y",
    log_tol: float = 2.0,
) -> None:
    """Conditionally sets an axis to logscale.

    If the logarithmic range of data spans at least log_tol orders of
    magnitude the specified axis is set to logscale.
    """
    log_diff = np.log10(np.max(data)) - np.log10(np.min(data))
    if log_diff > log_tol:
        if which_axis.lower() == "y":
            ax.set_yscale("log")
        elif which_axis.lower() == "x":
            ax.set_xscale("log")


def scale_figsize(rows: float, cols: float) -> tuple[float, float]:
    """Determines an appropriate figsize for a matrix of subplots.

    The new size is obtained by scaling the default figsize by
    a requested number (can be a fraction) of rows and columns.

    Parameters
    ----------
    rows : float
        Number of subplot rows to scale by.
    cols : float
        Number of subplot columns to scale by.

    Returns
    -------
    tuple of float
        A (width, height) tuple for the scaled figure size.
    """
    size = matplotlib.rcParams["figure.figsize"]
    new_size = (cols * size[0], rows * size[1])
    return new_size
