"""Plotting helpers for acfmperiod ACF results."""

from __future__ import annotations

from typing import List, Union

import matplotlib.pyplot as plt
import numpy as np

from .datatypes import AcfResult


def plotpanel(axis, xvals, yvals, plottype: str, color: str, **kwargs):
    """Draw one ACF panel on the provided axis.

    Parameters
    ----------
    axis
        Matplotlib axis that receives the plot.
    xvals
        Lag coordinates for the panel.
    yvals
        ACF/ACOVF values for the panel.
    plottype
        Plot style. ``"h"`` uses vertical-bar style; other values use
        ``axis.plot``.
    color
        Primary color used for the panel.
    **kwargs
        Extra keyword arguments passed to ``axis.plot`` when
        ``plottype != "h"``.
    """

    if plottype == "h":
        # Histogram-like bars are the default style in the reference R plots.
        axis.vlines(xvals, 0.0, yvals, colors=color, linewidth=1.2)
        axis.plot(xvals, yvals, "o", color=color, markersize=2.5)
    else:
        axis.plot(xvals, yvals, color=color, **kwargs)
    axis.axhline(0, color="black", linewidth=0.8)


def plotacfresult(
    x: AcfResult,
    plottype: str = "h",
    xlab: str = "Lag",
    ylab: str | None = None,
    ylim: tuple[float, float] | None = None,
    main: str | None = None,
    maxmfrow: int = 6,
    **kwargs,
) -> Union[plt.Figure, List[plt.Figure]]:
    """Plot an ACF result similarly to R ``acf`` and ``plot.robacf``.

    Parameters
    ----------
    x
        ACF/ACOVF result to plot.
    plottype
        Plot style. ``"h"`` for bar-like panels; otherwise line plotting.
    xlab
        X-axis label.
    ylab
        Y-axis label. If ``None``, a default is inferred from ``x``.
    ylim
        Y-axis limits. If ``None``, limits are inferred from data.
    main
        Panel title override. If ``None``, titles are inferred by series.
    maxmfrow
        Maximum number of rows/columns per plotted page.
    **kwargs
        Extra plotting options, including ``color``.

    Returns
    -------
    Union[plt.Figure, List[plt.Figure]]
        A single figure when one page is needed, otherwise a list of figures.

    Raises
    ------
    ValueError
        If ``x.lag`` is not a 3D tensor with at least one series.
    """

    if x.lag.ndim != 3:
        raise ValueError("x.lag must be a 3D array")

    nser = x.lag.shape[1]
    if nser < 1:
        raise ValueError("x.lag must have at least 1 column")

    if ylab is None:
        if x.robust:
            ylab = "Robust ACF" if x.type == "correlation" else "Robust ACF (cov)"
        else:
            ylab = "ACF" if x.type == "correlation" else "ACF (cov)"

    if x.snames is None:
        snames = [f"Series {index + 1}" for index in range(nser)]
    else:
        snames = [str(name) for name in x.snames]
        if len(snames) != nser:
            snames = [f"Series {index + 1}" for index in range(nser)]

    if ylim is None:
        if x.acf.size == 0:
            ylim = (-1.0, 1.0)
        else:
            minvalue = float(np.nanmin(x.acf))
            maxvalue = float(np.nanmax(x.acf))
            minvalue = min(minvalue, 0.0)
            maxvalue = max(maxvalue, 0.0)
            if minvalue == maxvalue:
                minvalue -= 1.0
                maxvalue += 1.0
            ylim = (minvalue, maxvalue)

    color = kwargs.pop("color", "C0")

    npages = 1
    nrows = nser
    if nser > maxmfrow:
        npages = int(np.ceil(nser / maxmfrow))
        nrows = int(np.ceil(nser / npages))

    figures: List[plt.Figure] = []
    for ipage in range(npages):
        for jpage in range(npages):
            figure, axes = plt.subplots(nrows, nrows, squeeze=False)
            iindices = np.arange(ipage * nrows, ipage * nrows + nrows)
            jindices = np.arange(jpage * nrows, jpage * nrows + nrows)

            for aindexi, i in enumerate(iindices):
                for aindexj, j in enumerate(jindices):
                    axis = axes[aindexi, aindexj]
                    if max(i, j) >= nser:
                        axis.axis("off")
                        continue

                    plotpanel(axis, x.lag[:, i, j], x.acf[:, i, j], plottype, color, **kwargs)
                    axis.set_ylim(*ylim)
                    axis.set_xlabel(xlab)
                    axis.set_ylabel(ylab if aindexj == 0 else "")

                    if main is not None:
                        title = main
                    elif i == j:
                        title = snames[i]
                    else:
                        title = f"{snames[i]} & {snames[j]}"
                    axis.set_title(title)

            figure.tight_layout()
            figures.append(figure)

    return figures[0] if len(figures) == 1 else figures


def plotrobacf(
    x: AcfResult,
    plottype: str = "h",
    xlab: str = "Lag",
    ylab: str | None = None,
    ylim: tuple[float, float] | None = None,
    main: str | None = None,
    maxmfrow: int = 6,
    **kwargs,
) -> Union[plt.Figure, List[plt.Figure]]:
    """Plot robust ACF results (kept for R API parity).

    Parameters
    ----------
    x
        Robust ACF/ACOVF result to plot.
    plottype
        Plot style. ``"h"`` for bar-like panels; otherwise line plotting.
    xlab
        X-axis label.
    ylab
        Y-axis label. If ``None``, a default is inferred.
    ylim
        Y-axis limits. If ``None``, limits are inferred from data.
    main
        Panel title override.
    maxmfrow
        Maximum number of rows/columns per plotted page.
    **kwargs
        Extra plotting options forwarded to :func:`plotacfresult`.

    Returns
    -------
    Union[plt.Figure, List[plt.Figure]]
        Figure output from :func:`plotacfresult`.
    """

    return plotacfresult(
        x=x,
        plottype=plottype,
        xlab=xlab,
        ylab=ylab,
        ylim=ylim,
        main=main,
        maxmfrow=maxmfrow,
        **kwargs,
    )
