from __future__ import annotations

import typing
from contextlib import suppress
from warnings import warn

import numpy as np
import pandas as pd

from .exceptions import PlotnineError, PlotnineWarning
from .facets import facet_grid, facet_null, facet_wrap
from .facets.facet_grid import parse_grid_facets
from .facets.facet_wrap import parse_wrap_facets
from .ggplot import ggplot
from .labels import labs
from .mapping.aes import ALL_AESTHETICS, SCALED_AESTHETICS, aes
from .scales import lims, scale_x_log10, scale_y_log10
from .themes import theme
from .utils import Registry, array_kind

if typing.TYPE_CHECKING:
    from typing import Any, Iterable, Literal

    from plotnine.typing import DataLike, Ggplot, TupleFloat2


def qplot(
    x: str | Iterable[Any] | range | None = None,
    y: str | Iterable[Any] | range | None = None,
    data: DataLike | None = None,
    facets: str = "",
    margins: bool | list[str] = False,
    geom: str | list[str] | tuple[str] = "auto",
    xlim: TupleFloat2 | None = None,
    ylim: TupleFloat2 | None = None,
    log: Literal["x", "y", "xy"] | None = None,
    main: str | None = None,
    xlab: str | None = None,
    ylab: str | None = None,
    asp: float | None = None,
    **kwargs: Any,
) -> Ggplot:
    """
    Quick plot

    Parameters
    ----------
    x : str | array_like
        x aesthetic
    y : str | array_like
        y aesthetic
    data : dataframe
        Data frame to use (optional). If not specified,
        will create one, extracting arrays from the
        current environment.
    geom : str | list
        *geom(s)* to do the drawing. If ``auto``, defaults
        to 'point' if ``x`` and ``y`` are specified or
        'histogram' if only ``x`` is specified.
    facets : str
        Facets
    margins : bool | list[str]
        variable names to compute margins for. True will compute
        all possible margins. Depends on the facetting.
    xlim : tuple
        x-axis limits
    ylim : tuple
        y-axis limits
    log : str in ``{'x', 'y', 'xy'}``
        Which (if any) variables to log transform.
    main : str
        Plot title
    xlab : str
        x-axis label
    ylab : str
        y-axis label
    asp : str | float
        The y/x aspect ratio.
    **kwargs : dict
        Arguments passed on to the geom.

    Returns
    -------
    p : ggplot
        ggplot object
    """
    from patsy.eval import EvalEnvironment

    # Extract all recognizable aesthetic mappings from the parameters
    # String values e.g  "I('red')", "I(4)" are not treated as mappings
    environment = EvalEnvironment.capture(1)
    aesthetics = {} if x is None else {"x": x}
    if y is not None:
        aesthetics["y"] = y

    def is_mapping(value: Any) -> bool:
        """
        Return True if value is not enclosed in I() function
        """
        with suppress(AttributeError):
            return not (value.startswith("I(") and value.endswith(")"))
        return True

    def I(value: Any) -> Any:
        return value

    I_env = EvalEnvironment([{"I": I}])

    for ae in kwargs.keys() & ALL_AESTHETICS:
        value = kwargs[ae]
        if is_mapping(value):
            aesthetics[ae] = value
        else:
            kwargs[ae] = I_env.eval(value)

    # List of geoms
    if isinstance(geom, str):
        geom = [geom]
    elif isinstance(geom, tuple):
        geom = list(geom)

    if data is None:
        data = pd.DataFrame()

    # Work out plot data, and modify aesthetics, if necessary
    def replace_auto(lst: list[str], str2: str) -> list[str]:
        """
        Replace all occurences of 'auto' in with str2
        """
        for i, value in enumerate(lst):
            if value == "auto":
                lst[i] = str2
        return lst

    if "auto" in geom:
        if "sample" in aesthetics:
            replace_auto(geom, "qq")
        elif y is None:
            # If x is discrete we choose geom_bar &
            # geom_histogram otherwise. But we need to
            # evaluate the mapping to find out the dtype
            env = environment.with_outer_namespace({"factor": pd.Categorical})

            if isinstance(aesthetics["x"], str):
                try:
                    x = env.eval(aesthetics["x"], inner_namespace=data)
                except Exception:
                    msg = "Could not evaluate aesthetic 'x={}'"
                    raise PlotnineError(msg.format(aesthetics["x"]))
            elif not hasattr(aesthetics["x"], "dtype"):
                x = np.asarray(aesthetics["x"])

            if array_kind.discrete(x):
                replace_auto(geom, "bar")
            else:
                replace_auto(geom, "histogram")

        else:
            if x is None:
                if isinstance(aesthetics["y"], typing.Sized):
                    aesthetics["x"] = range(len(aesthetics["y"]))
                    xlab = "range(len(y))"
                    ylab = "y"
                else:
                    # We could solve the issue in layer.compute_asthetics
                    # but it is not worth the extra complexity
                    raise PlotnineError("Cannot infer how long x should be.")
            replace_auto(geom, "point")

    p: Ggplot = ggplot(data, aes(**aesthetics), environment=environment)

    def get_facet_type(facets: str) -> Literal["grid", "wrap", "null"]:
        with suppress(PlotnineError):
            parse_grid_facets(facets)
            return "grid"

        with suppress(PlotnineError):
            parse_wrap_facets(facets)
            return "wrap"

        warn(
            "Could not determine the type of faceting, "
            "therefore no faceting.",
            PlotnineWarning,
        )
        return "null"

    if facets:
        facet_type = get_facet_type(facets)
        if facet_type == "grid":
            p += facet_grid(facets, margins=margins)
        elif facet_type == "wrap":
            p += facet_wrap(facets)
        else:
            p += facet_null()

    # Add geoms
    for g in geom:
        geom_name = f"geom_{g}"
        geom_klass = Registry[geom_name]
        stat_name = "stat_{}".format(geom_klass.DEFAULT_PARAMS["stat"])
        stat_klass = Registry[stat_name]
        # find params
        recognized = kwargs.keys() & (
            geom_klass.DEFAULT_PARAMS.keys()
            | geom_klass.aesthetics()
            | stat_klass.DEFAULT_PARAMS.keys()
            | stat_klass.aesthetics()
        )
        recognized = recognized - aesthetics.keys()
        params = {ae: kwargs[ae] for ae in recognized}
        p += geom_klass(**params)

    # pd.Series objects have name attributes. In a dataframe, the
    # series have the name of the column.
    labels = {}
    for ae in SCALED_AESTHETICS & kwargs.keys():
        with suppress(AttributeError):
            labels[ae] = kwargs[ae].name

    with suppress(AttributeError):
        labels["x"] = xlab if xlab is not None else x.name  # type: ignore

    with suppress(AttributeError):
        labels["y"] = ylab if ylab is not None else y.name  # type: ignore

    if main is not None:
        labels["title"] = main

    if log:
        if "x" in log:
            p += scale_x_log10()

        if "y" in log:
            p += scale_y_log10()

    if labels:
        p += labs(**labels)

    if asp:
        p += theme(aspect_ratio=asp)

    if xlim:
        p += lims(x=xlim)

    if ylim:
        p += lims(y=ylim)

    return p
