"""ACF and ACOVF estimators reconstructed from periodogram quantities."""

from __future__ import annotations

from typing import Optional

import numpy as np
from numpy.typing import ArrayLike, NDArray

from .datatypes import AcfResult, NAAction
from .periodograms import crossperiodogram, mcrossperiodogram, mperioreg, perioreg
from .validators import applynaaction, matchtype, preparetwodim


def gmat(n: int) -> NDArray[np.complex128]:
    """Build the Fourier matrix used in covariance reconstruction.

    Parameters
    ----------
    n
        Matrix dimension and series length.

    Returns
    -------
    NDArray[np.complex128]
        Unitary Fourier matrix with shape ``(n, n)``.

    Raises
    ------
    ValueError
        If ``n < 1``.
    """

    if n < 1:
        raise ValueError("n must be >= 1")
    idx = np.arange(n, dtype=float)
    matrixg = np.exp(1j * 2 * np.pi * np.outer(idx, idx) / n)
    return (1.0 / np.sqrt(n)) * matrixg


def covfromspectrum(spectrum: NDArray[np.complex128], factor: float) -> NDArray[np.float64]:
    """Recover a covariance matrix from diagonal spectral values.

    Parameters
    ----------
    spectrum
        Complex spectrum values used as the diagonal of a frequency-domain
        covariance operator.
    factor
        Scaling factor applied to the reconstructed covariance matrix.

    Returns
    -------
    NDArray[np.float64]
        Real-valued covariance matrix.
    """

    matrixg = gmat(spectrum.size)
    # Equivalent to Conj(t(G)) %*% D %*% G in the R implementation.
    weighted = spectrum[:, None] * matrixg
    covmatrix = factor * np.real(np.conjugate(matrixg.T) @ weighted)
    return covmatrix.astype(float, copy=False)


def buildlagcube(nser: int, lagmax: int, xfreq: float = 1.0) -> NDArray[np.float64]:
    """Build lag tensor layout compatible with acf-like outputs.

    Parameters
    ----------
    nser
        Number of series.
    lagmax
        Number of returned lags.
    xfreq
        Sampling frequency used to scale lag units.

    Returns
    -------
    NDArray[np.float64]
        Lag tensor with shape ``(lagmax, nser, nser)``.
    """

    lagsign = np.ones((nser, nser), dtype=float)
    lagsign[np.tril_indices(nser, k=-1)] = -1.0
    if lagmax == 0:
        return np.empty((0, nser, nser), dtype=float)
    return np.multiply.outer(np.arange(lagmax, dtype=float), lagsign / xfreq)


def resolvelagmax(samplet: int, nser: int, lagmax: Optional[int]) -> int:
    """Resolve lag maximum with defaults/bounds matching the R package.

    Parameters
    ----------
    samplet
        Number of observations.
    nser
        Number of series.
    lagmax
        User-specified lag maximum or ``None`` for automatic selection.

    Returns
    -------
    int
        Effective lag maximum bounded by ``samplet - 1``.

    Raises
    ------
    ValueError
        If the resolved lag maximum is negative.
    """

    if lagmax is None:
        lagmax = int(np.floor(10 * (np.log10(samplet) - np.log10(nser))))
    lagmax = int(min(lagmax, samplet - 1))
    if lagmax < 0:
        raise ValueError("lagmax must be at least 0")
    return lagmax


def normalizetocorrelation(acfvals: NDArray[np.float64], lagmax: int, nser: int) -> None:
    """Normalize covariance estimates in-place to correlation estimates.

    Parameters
    ----------
    acfvals
        Covariance estimates with shape ``(lagmax, nser, nser)``.
    lagmax
        Number of lags represented in ``acfvals``.
    nser
        Number of series.

    Returns
    -------
    None
        The array is modified in-place.
    """

    if lagmax == 0:
        return
    for i in range(nser):
        for j in range(nser):
            if i != j:
                denom = np.sqrt(acfvals[0, i, i]) * np.sqrt(acfvals[0, j, j])
                acfvals[:, i, j] = acfvals[:, i, j] / denom
    for i in range(nser):
        acfvals[:, i, i] = acfvals[:, i, i] / acfvals[0, i, i]


def acfcore(
    x: ArrayLike,
    lagmax: Optional[int],
    typevalue: str,
    doplot: bool,
    naaction: NAAction,
    demean: bool,
    robust: bool,
    **plotkwargs,
) -> AcfResult:
    """Shared core for classical and robust ACF/ACOVF estimators.

    Parameters
    ----------
    x
        Input vector or matrix-like object of observations.
    lagmax
        Number of lags to compute. If ``None``, use the package default rule.
    typevalue
        Correlation/covariance selector accepted by :func:`matchtype`.
    doplot
        If ``True``, render the result via :meth:`AcfResult.plot`.
    naaction
        Missing-data policy. See :func:`applynaaction`.
    demean
        If ``True``, subtract the column mean before estimation.
    robust
        If ``True``, use robust periodogram/cross-periodogram estimators.
    **plotkwargs
        Keyword arguments forwarded to plotting helpers when ``doplot=True``.

    Returns
    -------
    AcfResult
        ACF/ACOVF tensor, lag tensor, and metadata.

    Raises
    ------
    ValueError
        If inputs are invalid, contain NaN values under ``naaction='fail'``,
        or request an invalid ``typevalue``.
    """

    acftype = matchtype(typevalue)
    xarr, snames = preparetwodim(x)
    xarr = applynaaction(xarr, naaction)

    if not np.issubdtype(xarr.dtype, np.number):
        raise ValueError("x must be numeric")

    samplet, nser = xarr.shape
    lagmaxvalue = resolvelagmax(samplet, nser, lagmax)

    if demean:
        xarr = xarr - np.nanmean(xarr, axis=0, keepdims=True)

    acfvals = np.ones((lagmaxvalue, nser, nser), dtype=float)
    periodfunc = mperioreg if robust else perioreg
    crossfunc = mcrossperiodogram if robust else crossperiodogram

    if nser == 1:
        spectrum = np.concatenate([[0.0], periodfunc(xarr[:, 0])]).astype(complex)
        covmatrix = covfromspectrum(spectrum, factor=2 * np.pi)
        if lagmaxvalue > 0:
            acfvals[:, 0, 0] = covmatrix[0, :lagmaxvalue]
            if acftype == "correlation":
                acfvals = acfvals / acfvals[0, 0, 0]
    else:
        # Keep iteration order aligned with the original R implementation.
        for i in range(nser):
            jrange = range(nser) if robust else range(i, nser)
            for j in jrange:
                if i == j:
                    spectrum = np.concatenate([[0.0], periodfunc(xarr[:, i])]).astype(complex)
                    covmatrix = covfromspectrum(spectrum, factor=2 * np.pi)
                    if lagmaxvalue > 0:
                        acfvals[:, i, j] = covmatrix[0, :lagmaxvalue]
                else:
                    cross = crossfunc(xarr[:, i], xarr[:, j])
                    covxy = covfromspectrum(cross.crossperiodxy, factor=0.5)
                    covyx = covfromspectrum(cross.crossperiodyx, factor=0.5)
                    if lagmaxvalue > 0:
                        acfvals[:, i, j] = covxy[0, :lagmaxvalue]
                        acfvals[:, j, i] = covyx[0, :lagmaxvalue]

        if acftype == "correlation":
            normalizetocorrelation(acfvals, lagmaxvalue, nser)

    result = AcfResult(
        acf=acfvals,
        type=acftype,
        nused=samplet,
        lag=buildlagcube(nser, lagmaxvalue),
        series="x",
        snames=snames,
        robust=robust,
    )
    if doplot:
        result.plot(**plotkwargs)
    return result


def peracf(
    x: ArrayLike,
    lagmax: Optional[int] = None,
    typevalue: str = "correlation",
    doplot: bool = True,
    naaction: NAAction = "fail",
    demean: bool = True,
    **plotkwargs,
) -> AcfResult:
    """Estimate classical ACF/ACOVF from periodogram quantities.

    Parameters
    ----------
    x
        Input vector or matrix-like object.
    lagmax
        Number of lags to compute. If ``None``, a default based on sample size
        and dimensionality is used.
    typevalue
        ``"correlation"`` or ``"covariance"`` (prefixes allowed).
    doplot
        If ``True``, plot the resulting estimate.
    naaction
        Missing-data policy (``"fail"``, ``"pass"``, or callable).
    demean
        If ``True``, subtract column means before estimation.
    **plotkwargs
        Plotting options forwarded to :func:`acfmperiod.plotting.plotacfresult`.

    Returns
    -------
    AcfResult
        Classical periodogram-based ACF/ACOVF result.
    """

    return acfcore(
        x=x,
        lagmax=lagmax,
        typevalue=typevalue,
        doplot=doplot,
        naaction=naaction,
        demean=demean,
        robust=False,
        **plotkwargs,
    )


def mperacf(
    x: ArrayLike,
    lagmax: Optional[int] = None,
    typevalue: str = "correlation",
    doplot: bool = True,
    naaction: NAAction = "fail",
    demean: bool = True,
    **plotkwargs,
) -> AcfResult:
    """Estimate robust ACF/ACOVF from M-periodogram quantities.

    Parameters
    ----------
    x
        Input vector or matrix-like object.
    lagmax
        Number of lags to compute. If ``None``, a default based on sample size
        and dimensionality is used.
    typevalue
        ``"correlation"`` or ``"covariance"`` (prefixes allowed).
    doplot
        If ``True``, plot the resulting estimate.
    naaction
        Missing-data policy (``"fail"``, ``"pass"``, or callable).
    demean
        If ``True``, subtract column means before estimation.
    **plotkwargs
        Plotting options forwarded to :func:`acfmperiod.plotting.plotacfresult`.

    Returns
    -------
    AcfResult
        Robust periodogram-based ACF/ACOVF result.
    """

    return acfcore(
        x=x,
        lagmax=lagmax,
        typevalue=typevalue,
        doplot=doplot,
        naaction=naaction,
        demean=demean,
        robust=True,
        **plotkwargs,
    )


def covcorper(x: ArrayLike, typevalue: str = "correlation") -> NDArray[np.float64]:
    """Return lag-0 covariance/correlation matrix from :func:`peracf`.

    Parameters
    ----------
    x
        Input vector or matrix-like object.
    typevalue
        ``"correlation"`` or ``"covariance"`` (prefixes allowed).

    Returns
    -------
    NDArray[np.float64]
        Lag-0 covariance or correlation matrix.
    """

    matrix = peracf(x, lagmax=1, doplot=False, typevalue=typevalue).acf
    return matrix[0, :, :]


def covcormper(x: ArrayLike, typevalue: str = "correlation") -> NDArray[np.float64]:
    """Return lag-0 covariance/correlation matrix from :func:`mperacf`.

    Parameters
    ----------
    x
        Input vector or matrix-like object.
    typevalue
        ``"correlation"`` or ``"covariance"`` (prefixes allowed).

    Returns
    -------
    NDArray[np.float64]
        Lag-0 robust covariance or correlation matrix.
    """

    matrix = mperacf(x, lagmax=1, doplot=False, typevalue=typevalue).acf
    return matrix[0, :, :]
