"""Shared data containers and typing aliases for acfmperiod."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Optional, Sequence, Union

import numpy as np
from numpy.typing import NDArray


# This type mirrors the original R argument style:
# keyword-like values or a custom callable.
NAAction = Union[str, Callable[[NDArray[np.float64]], NDArray[np.float64]]]


@dataclass
class CrossPeriodogramResult:
    """Container for directed cross-periodogram estimates.

    Attributes
    ----------
    crossperiodxy
        Frequency-domain estimate for direction ``x -> y``.
    crossperiodyx
        Frequency-domain estimate for direction ``y -> x``.
    """

    crossperiodxy: NDArray[np.complex128]
    crossperiodyx: NDArray[np.complex128]


@dataclass
class AcfResult:
    """Result object returned by periodogram-based ACF estimators.

    Attributes
    ----------
    acf
        Estimated ACF/ACOVF tensor with shape ``(lagmax, nser, nser)``.
    type
        ``"correlation"`` or ``"covariance"``.
    nused
        Number of observations used in estimation.
    lag
        Lag tensor aligned with ``acf``.
    series
        Name of the source object, kept for API parity.
    snames
        Optional sequence of series names.
    robust
        Whether robust M-periodogram estimation was used.
    """

    acf: NDArray[np.float64]
    type: str
    nused: int
    lag: NDArray[np.float64]
    series: str = "x"
    snames: Optional[Sequence[str]] = None
    robust: bool = False

    def plot(self, **kwargs):
        """Render the result using package plotting helpers.

        Parameters
        ----------
        **kwargs
            Keyword arguments forwarded to
            :func:`acfmperiod.plotting.plotacfresult`.

        Returns
        -------
        matplotlib.figure.Figure or list[matplotlib.figure.Figure]
            Generated figure object(s).
        """

        from .plotting import plotacfresult

        return plotacfresult(self, **kwargs)
