"""Periodogram and cross-periodogram estimators."""

from __future__ import annotations

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

from .datatypes import CrossPeriodogramResult
from .regressions import fitols, fitrlm
from .validators import prepareonedim


def mirrorhalf(values: NDArray, n: int) -> NDArray:
    """Mirror positive-frequency values to match package layout.

    Parameters
    ----------
    values
        Values computed on positive Fourier frequencies (length ``n // 2``).
    n
        Original sample size.

    Returns
    -------
    NDArray
        Mirrored vector with length ``n - 1`` matching the layout used by the
        original R package.
    """

    if n % 2:
        return np.concatenate([values, values[::-1]])
    return np.concatenate([values, values[::-1][1:]])


def harmoniccoeffhalf(series: NDArray[np.float64], robust: bool) -> NDArray[np.complex128]:
    """Estimate complex harmonic coefficients at positive frequencies.

    Parameters
    ----------
    series
        One-dimensional numeric series.
    robust
        If ``True``, use Huber M-regression; otherwise use ordinary least
        squares.

    Returns
    -------
    NDArray[np.complex128]
        Complex harmonic coefficients for positive frequencies.
    """

    n = len(series)
    half = n // 2
    coeffs = np.empty(half, dtype=np.complex128)
    tindex = np.arange(1, n + 1, dtype=float)
    fitfunc = fitrlm if robust else fitols

    for j in range(1, half + 1):
        omega = 2 * np.pi * j / n
        cospart = np.cos(omega * tindex)
        if j != (n / 2):
            sinpart = np.sin(omega * tindex)
            matrixx = np.column_stack([cospart, sinpart])
            coef = fitfunc(series, matrixx)
            coeffs[j - 1] = complex(coef[0], -coef[1])
        else:
            matrixx = cospart[:, None]
            coef = fitfunc(series, matrixx)
            coeffs[j - 1] = complex(coef[0], 0.0)
    return coeffs


def periodogramcore(series: ArrayLike, robust: bool) -> NDArray[np.float64]:
    """Shared implementation for classical and robust periodograms.

    Parameters
    ----------
    series
        Input numeric series.
    robust
        If ``True``, use robust harmonic regression; otherwise use classical
        least-squares regression.

    Returns
    -------
    NDArray[np.float64]
        Periodogram values mirrored to the package layout (length ``n - 1``).

    Raises
    ------
    ValueError
        If ``series`` is invalid according to :func:`prepareonedim`.
    """

    y = prepareonedim(series)
    n = len(y)
    half = n // 2
    coeffhalf = harmoniccoeffhalf(y, robust=robust)
    periorhalf = np.empty(half, dtype=float)

    for j in range(1, half + 1):
        if j != (n / 2):
            fft = np.sqrt(n / (8 * np.pi)) * coeffhalf[j - 1]
        else:
            fft = np.sqrt(n / (2 * np.pi)) * complex(coeffhalf[j - 1].real, 0.0)
        periorhalf[j - 1] = abs(fft) ** 2
    return mirrorhalf(periorhalf, n)


def periodogramauxcore(series: ArrayLike, robust: bool) -> NDArray[np.complex128]:
    """Return mirrored complex harmonic coefficients.

    Parameters
    ----------
    series
        Input numeric series.
    robust
        If ``True``, estimate robust coefficients; otherwise estimate classical
        coefficients.

    Returns
    -------
    NDArray[np.complex128]
        Mirrored complex harmonic coefficients used by cross-periodogram
        estimators.
    """

    y = prepareonedim(series)
    n = len(y)
    coeffhalf = harmoniccoeffhalf(y, robust=robust)
    return mirrorhalf(coeffhalf, n)


def crossperiodogramcore(
    series1: ArrayLike,
    series2: ArrayLike,
    robust: bool,
) -> CrossPeriodogramResult:
    """Shared implementation for classical and robust cross-periodograms.

    Parameters
    ----------
    series1
        First input series.
    series2
        Second input series. Must have the same length as ``series1``.
    robust
        If ``True``, use robust harmonic regression; otherwise use classical
        least-squares regression.

    Returns
    -------
    CrossPeriodogramResult
        Directed cross-periodogram estimates for ``(series1, series2)`` and
        ``(series2, series1)``.

    Raises
    ------
    ValueError
        If inputs are invalid one-dimensional series or have different lengths.
    """

    s1 = prepareonedim(series1)
    s2 = prepareonedim(series2)
    if len(s1) != len(s2):
        raise ValueError("series1 and series2 must have identical length")

    n = len(s1)
    period1 = periodogramauxcore(s1, robust=robust)
    period2 = periodogramauxcore(s2, robust=robust)

    alpha1 = period1.real
    beta1 = period1.imag
    alpha2 = period2.real
    beta2 = period2.imag

    # This algebra mirrors the original R implementation.
    realpart = alpha1 * alpha2 + beta1 * beta2
    imagdelta = alpha1 * beta2 - alpha2 * beta1
    crossxy = (n / 2.0) * (realpart - 1j * imagdelta)
    crossyx = (n / 2.0) * (realpart + 1j * imagdelta)

    prefix = 2 * n * alpha1[-1] * alpha2[-1]
    return CrossPeriodogramResult(
        crossperiodxy=np.concatenate([[complex(prefix, 0.0)], crossxy]),
        crossperiodyx=np.concatenate([[complex(prefix, 0.0)], crossyx]),
    )


def perioreg(series: ArrayLike) -> NDArray[np.float64]:
    """Compute the non-robust periodogram via harmonic regression.

    Parameters
    ----------
    series
        Input one-dimensional numeric series.

    Returns
    -------
    NDArray[np.float64]
        Mirrored periodogram values with length ``n - 1``.
    """

    return periodogramcore(series, robust=False)


def mperioreg(series: ArrayLike) -> NDArray[np.float64]:
    """Compute the robust M-periodogram via Huber M-regression.

    Parameters
    ----------
    series
        Input one-dimensional numeric series.

    Returns
    -------
    NDArray[np.float64]
        Mirrored robust periodogram values with length ``n - 1``.
    """

    return periodogramcore(series, robust=True)


def crossperiodogram(series1: ArrayLike, series2: ArrayLike) -> CrossPeriodogramResult:
    """Compute the non-robust cross-periodogram.

    Parameters
    ----------
    series1
        First input series.
    series2
        Second input series with the same length as ``series1``.

    Returns
    -------
    CrossPeriodogramResult
        Directed non-robust cross-periodogram pair.
    """

    return crossperiodogramcore(series1, series2, robust=False)


def mcrossperiodogram(series1: ArrayLike, series2: ArrayLike) -> CrossPeriodogramResult:
    """Compute the robust M-cross-periodogram.

    Parameters
    ----------
    series1
        First input series.
    series2
        Second input series with the same length as ``series1``.

    Returns
    -------
    CrossPeriodogramResult
        Directed robust cross-periodogram pair.
    """

    return crossperiodogramcore(series1, series2, robust=True)
