1"""Matplotlib kdeplot."""
2import numpy as np
3from matplotlib import pyplot as plt
4from matplotlib import _pylab_helpers
5import matplotlib.ticker as mticker
6
7
8from ...plot_utils import _scale_fig_size
9from . import backend_kwarg_defaults, backend_show, create_axes_grid, matplotlib_kwarg_dealiaser
10
11
12def plot_kde(
13    density,
14    lower,
15    upper,
16    density_q,
17    xmin,
18    xmax,
19    ymin,
20    ymax,
21    gridsize,
22    values,
23    values2,
24    rug,
25    label,
26    quantiles,
27    rotated,
28    contour,
29    fill_last,
30    figsize,
31    textsize,
32    plot_kwargs,
33    fill_kwargs,
34    rug_kwargs,
35    contour_kwargs,
36    contourf_kwargs,
37    pcolormesh_kwargs,
38    is_circular,
39    ax,
40    legend,
41    backend_kwargs,
42    show,
43    return_glyph,  # pylint: disable=unused-argument
44):
45    """Matplotlib kde plot."""
46    if backend_kwargs is None:
47        backend_kwargs = {}
48
49    backend_kwargs = {
50        **backend_kwarg_defaults(),
51        **backend_kwargs,
52    }
53
54    figsize, *_, xt_labelsize, linewidth, markersize = _scale_fig_size(figsize, textsize)
55
56    backend_kwargs.setdefault("figsize", figsize)
57    backend_kwargs["squeeze"] = True
58    backend_kwargs.setdefault("subplot_kw", {})
59    backend_kwargs["subplot_kw"].setdefault("polar", is_circular)
60
61    if ax is None:
62        fig_manager = _pylab_helpers.Gcf.get_active()
63        if fig_manager is not None:
64            ax = fig_manager.canvas.figure.gca()
65        else:
66            _, ax = create_axes_grid(
67                1,
68                backend_kwargs=backend_kwargs,
69            )
70
71    if values2 is None:
72        plot_kwargs = matplotlib_kwarg_dealiaser(plot_kwargs, "plot")
73        plot_kwargs.setdefault("color", "C0")
74
75        default_color = plot_kwargs.get("color")
76
77        fill_kwargs = matplotlib_kwarg_dealiaser(fill_kwargs, "hexbin")
78        fill_kwargs.setdefault("color", default_color)
79
80        rug_kwargs = matplotlib_kwarg_dealiaser(rug_kwargs, "plot")
81        rug_kwargs.setdefault("marker", "_" if rotated else "|")
82        rug_kwargs.setdefault("linestyle", "None")
83        rug_kwargs.setdefault("color", default_color)
84        rug_kwargs.setdefault("space", 0.2)
85
86        plot_kwargs.setdefault("linewidth", linewidth)
87        rug_kwargs.setdefault("markersize", 2 * markersize)
88
89        rug_space = max(density) * rug_kwargs.pop("space")
90
91        if is_circular:
92
93            if is_circular == "radians":
94                labels = [
95                    "0",
96                    f"{np.pi/4:.2f}",
97                    f"{np.pi/2:.2f}",
98                    f"{3*np.pi/4:.2f}",
99                    f"{np.pi:.2f}",
100                    f"{-3*np.pi/4:.2f}",
101                    f"{-np.pi/2:.2f}",
102                    f"{-np.pi/4:.2f}",
103                ]
104
105                ticks_loc = ax.get_xticks()
106                ax.xaxis.set_major_locator(mticker.FixedLocator(ticks_loc))
107                ax.set_xticklabels(labels)
108
109            x = np.linspace(-np.pi, np.pi, len(density))
110            ax.set_yticklabels([])
111
112        else:
113            x = np.linspace(lower, upper, len(density))
114
115        fill_func = ax.fill_between
116        fill_x, fill_y = x, density
117        if rotated:
118            x, density = density, x
119            fill_func = ax.fill_betweenx
120
121        ax.tick_params(labelsize=xt_labelsize)
122
123        if rotated:
124            ax.set_xlim(0, auto=True)
125            rug_x, rug_y = np.zeros_like(values) - rug_space, values
126        else:
127            ax.set_ylim(0, auto=True)
128            rug_x, rug_y = values, np.zeros_like(values) - rug_space
129
130        if rug:
131            ax.plot(rug_x, rug_y, **rug_kwargs)
132
133        if quantiles is not None:
134            fill_kwargs.setdefault("alpha", 0.75)
135
136            idx = [np.sum(density_q < quant) for quant in quantiles]
137
138            fill_func(
139                fill_x,
140                fill_y,
141                where=np.isin(fill_x, fill_x[idx], invert=True, assume_unique=True),
142                **fill_kwargs,
143            )
144        else:
145            fill_kwargs.setdefault("alpha", 0)
146            if fill_kwargs.get("alpha") == 0:
147                label = plot_kwargs.setdefault("label", label)
148                ax.plot(x, density, **plot_kwargs)
149                fill_func(fill_x, fill_y, **fill_kwargs)
150            else:
151                label = fill_kwargs.setdefault("label", label)
152                ax.plot(x, density, **plot_kwargs)
153                fill_func(fill_x, fill_y, **fill_kwargs)
154        if legend and label:
155            ax.legend()
156    else:
157        contour_kwargs = matplotlib_kwarg_dealiaser(contour_kwargs, "contour")
158        contour_kwargs.setdefault("colors", "0.5")
159        contourf_kwargs = matplotlib_kwarg_dealiaser(contourf_kwargs, "contour")
160        pcolormesh_kwargs = matplotlib_kwarg_dealiaser(pcolormesh_kwargs, "pcolormesh")
161
162        g_s = complex(gridsize[0])
163        x_x, y_y = np.mgrid[xmin:xmax:g_s, ymin:ymax:g_s]
164
165        ax.grid(False)
166        ax.set_xlim(xmin, xmax)
167        ax.set_ylim(ymin, ymax)
168        if contour:
169            qcfs = ax.contourf(x_x, y_y, density, antialiased=True, **contourf_kwargs)
170            qcs = ax.contour(x_x, y_y, density, **contour_kwargs)
171            if not fill_last:
172                qcfs.collections[0].set_alpha(0)
173                qcs.collections[0].set_alpha(0)
174        else:
175            ax.pcolormesh(x_x, y_y, density, **pcolormesh_kwargs)
176
177    if backend_show(show):
178        plt.show()
179
180    return ax
181