"""
Fast polyphase resampling with multi-architecture SIMD support.

This module provides high-performance upsampling using polyphase filtering,
with automatic CPU detection and SIMD optimization.

Supported architectures:
- x86_64: AVX-512, AVX2+FMA, AVX, SSE4.1, SSE2
- ARM64: NEON (Apple Silicon, AWS Graviton, etc.)
- All: Optimized scalar fallback

Two memory layouts are supported:
- Standard: (time, channels) - use upsample()
- Transposed: (channels, time) - use upsample_transposed()
"""

import numpy as np

# Import from the SIMD extension
from sgnl_cpu_interp._core import correlate as _correlate_core
from sgnl_cpu_interp._core import get_available_impls as _get_available_impls
from sgnl_cpu_interp._core import get_cpu_features as _get_cpu_features
from sgnl_cpu_interp._core import get_impl_name as _get_impl_name
from sgnl_cpu_interp._core import select_impl as _select_impl
from sgnl_cpu_interp._core import upsample as _upsample_core
from sgnl_cpu_interp._core import upsample_transposed as _upsample_transposed_core
from sgnl_cpu_interp._version import version as __version__


def get_simd_info():
    """
    Get information about the current SIMD implementation.

    Returns
    -------
    dict
        Dictionary containing:
        - 'implementation': Name of current implementation (e.g., 'AVX2+FMA', 'NEON')
        - 'available': List of all available implementations
        - 'cpu_features': Detected CPU SIMD features

    Examples
    --------
    >>> from sgnl_cpu_interp import get_simd_info
    >>> info = get_simd_info()
    >>> print(f"Using: {info['implementation']}")
    Using: NEON
    """
    return {
        "implementation": _get_impl_name(),
        "available": _get_available_impls(),
        "cpu_features": _get_cpu_features(),
    }


def set_implementation(name):
    """
    Manually select a SIMD implementation.

    This is primarily useful for testing and benchmarking different
    implementations. In normal use, the best implementation is
    automatically selected at module load time.

    Can also be set via the SGNL_CPU_IMPL environment variable.

    Parameters
    ----------
    name : str
        Implementation name. Use get_simd_info()['available'] to see options.

    Raises
    ------
    ValueError
        If implementation name is not recognized
    RuntimeError
        If CPU features required for the implementation are not available

    Examples
    --------
    >>> from sgnl_cpu_interp import set_implementation, get_simd_info
    >>> set_implementation('Scalar')
    >>> print(get_simd_info()['implementation'])
    Scalar
    """
    _select_impl(name)


def upsample(data, factor=2, half_length=8):
    """
    Upsample multichannel data using polyphase filtering.

    This function performs high-quality upsampling using a Lanczos-windowed
    sinc interpolation kernel. It's optimized for processing many channels
    simultaneously using SIMD instructions.

    Parameters
    ----------
    data : ndarray
        Input array of shape (n_samples, n_channels). Will be converted
        to float32 if necessary.
    factor : int, optional
        Upsampling factor (default: 2). Must be >= 2.
    half_length : int, optional
        Half-length of the sinc kernel (default: 8). Larger values provide
        better quality but are slower. The total kernel length will be
        2 * half_length + 1.

    Returns
    -------
    output : ndarray
        Upsampled array of shape ((n_samples - kernel_len + 1) * factor, n_channels)
        where kernel_len = 2 * half_length + 1. Output is float32.

    Notes
    -----
    The function uses a Lanczos-windowed sinc kernel for interpolation:

        h(x) = sinc(x/factor) * sinc(x/half_length)

    Edge samples are lost due to the convolution (kernel_len - 1 samples total).

    Examples
    --------
    >>> import numpy as np
    >>> from sgnl_cpu_interp import upsample
    >>>
    >>> # Create test signal with 128 channels
    >>> n_samples, n_channels = 1024, 128
    >>> data = np.random.randn(n_samples, n_channels).astype(np.float32)
    >>>
    >>> # Upsample by factor of 2
    >>> upsampled = upsample(data, factor=2)
    >>> print(upsampled.shape)  # (2016, 128) - lost 16 edge samples

    >>> # Upsample by factor of 4 with longer kernel
    >>> upsampled = upsample(data, factor=4, half_length=16)
    >>> print(upsampled.shape)  # (3972, 128)
    """
    # Validate inputs
    if factor < 2:
        raise ValueError(f"factor must be >= 2, got {factor}")
    if half_length < 1:
        raise ValueError(f"half_length must be >= 1, got {half_length}")

    # Convert to numpy array if needed
    data = np.asarray(data)

    # Check dimensions
    if data.ndim == 1:
        # Single channel - add channel dimension
        data = data[:, np.newaxis]
        squeeze_output = True
    elif data.ndim == 2:
        squeeze_output = False
    else:
        raise ValueError(f"data must be 1D or 2D, got shape {data.shape}")

    # Convert to float32 and ensure C-contiguous
    if data.dtype != np.float32:
        data = data.astype(np.float32, order="C")
    elif not data.flags["C_CONTIGUOUS"]:
        data = np.ascontiguousarray(data)

    # Check minimum size
    kernel_len = 2 * half_length + 1
    if data.shape[0] < kernel_len:
        raise ValueError(
            f"data must have at least {kernel_len} samples (2*half_length+1), "
            f"got {data.shape[0]}"
        )

    # Call SIMD-optimized C extension
    result = _upsample_core(data, factor=factor, half_length=half_length)

    # Remove channel dimension if input was 1D
    if squeeze_output:
        result = result.squeeze(axis=1)

    return result


def upsample_transposed(data, factor=2, half_length=8):
    """
    Upsample multichannel data using polyphase filtering (transposed layout).

    This function is identical to `upsample()` but expects input in
    transposed layout: (n_channels, n_samples) instead of (n_samples, n_channels).

    This layout can be more efficient when channels are processed independently,
    as each channel's samples are contiguous in memory.

    Parameters
    ----------
    data : ndarray
        Input array of shape (n_channels, n_samples). Will be converted
        to float32 if necessary.
    factor : int, optional
        Upsampling factor (default: 2). Must be >= 2.
    half_length : int, optional
        Half-length of the sinc kernel (default: 8). Larger values provide
        better quality but are slower. The total kernel length will be
        2 * half_length + 1.

    Returns
    -------
    output : ndarray
        Upsampled array of shape (n_channels, (n_samples - kernel_len + 1) * factor)
        where kernel_len = 2 * half_length + 1. Output is float32.

    Examples
    --------
    >>> import numpy as np
    >>> from sgnl_cpu_interp import upsample_transposed
    >>>
    >>> # Create test signal with 128 channels (transposed layout)
    >>> n_channels, n_samples = 128, 1024
    >>> data = np.random.randn(n_channels, n_samples).astype(np.float32)
    >>>
    >>> # Upsample by factor of 2
    >>> upsampled = upsample_transposed(data, factor=2)
    >>> print(upsampled.shape)  # (128, 2016)
    """
    # Validate inputs
    if factor < 2:
        raise ValueError(f"factor must be >= 2, got {factor}")
    if half_length < 1:
        raise ValueError(f"half_length must be >= 1, got {half_length}")

    # Convert to numpy array if needed
    data = np.asarray(data)

    # Check dimensions
    if data.ndim != 2:
        raise ValueError(
            f"data must be 2D for transposed layout, got shape {data.shape}"
        )

    # Convert to float32 and ensure C-contiguous
    if data.dtype != np.float32:
        data = data.astype(np.float32, order="C")
    elif not data.flags["C_CONTIGUOUS"]:
        data = np.ascontiguousarray(data)

    # Check minimum size
    kernel_len = 2 * half_length + 1
    if data.shape[1] < kernel_len:
        raise ValueError(
            f"data must have at least {kernel_len} samples (2*half_length+1), "
            f"got {data.shape[1]}"
        )

    # Call SIMD-optimized C extension
    result = _upsample_transposed_core(data, factor=factor, half_length=half_length)

    return result


def correlate(data, filters, groups=None):
    """
    Grouped 1D correlation matching PyTorch conv1d with groups parameter.

    This is equivalent to torch.nn.functional.conv1d with the specified groups,
    but optimized for CPU with SIMD instructions (NEON, AVX2, etc.).

    Parameters
    ----------
    data : ndarray
        Input array of shape (n_channels, samples). Will be converted to float32
        if necessary.
    filters : ndarray
        Filter array of shape (n_filters, filter_length). Will be converted to
        float32 if necessary.
    groups : int, optional
        Number of groups for grouped convolution. Default is n_channels
        (depthwise convolution). Must divide both n_channels and n_filters evenly.

    Returns
    -------
    output : ndarray
        Output array of shape (n_filters, samples - filter_length + 1).
        Output is float32.

    Notes
    -----
    This function computes grouped correlation matching PyTorch's conv1d:

    With groups=g, the input channels and filters are divided into g groups.
    Each group of filters only operates on its corresponding group of input
    channels. This requires n_channels % groups == 0 and n_filters % groups == 0.

    Common cases:
    - groups=1: Standard convolution. All filters applied to all input channels.
      Example: data (1, 1000), filters (80, 64) -> output (80, 937)
    - groups=n_channels=n_filters: Depthwise convolution. Each filter applied
      to its corresponding input channel.
      Example: data (80, 1000), filters (80, 64) -> output (80, 937)

    This is correlation (not convolution), which matches PyTorch's conv1d
    behavior where the filter is NOT flipped.

    Examples
    --------
    >>> import numpy as np
    >>> from sgnl_cpu_interp import correlate
    >>>
    >>> # Depthwise: 64 channels, each with its own filter
    >>> data = np.random.randn(64, 16384).astype(np.float32)
    >>> filters = np.random.randn(64, 1024).astype(np.float32)
    >>> output = correlate(data, filters)  # groups=64 (default)
    >>> print(output.shape)  # (64, 15361)
    >>>
    >>> # Standard: 1 channel, 80 filters
    >>> data = np.random.randn(1, 16384).astype(np.float32)
    >>> filters = np.random.randn(80, 1024).astype(np.float32)
    >>> output = correlate(data, filters, groups=1)
    >>> print(output.shape)  # (80, 15361)
    """
    # Convert to numpy arrays if needed
    data = np.asarray(data)
    filters = np.asarray(filters)

    # Check dimensions
    if data.ndim != 2:
        raise ValueError(f"data must be 2D (channels, samples), got shape {data.shape}")
    if filters.ndim != 2:
        raise ValueError(
            f"filters must be 2D (n_filters, filter_length), got shape {filters.shape}"
        )

    n_channels = data.shape[0]
    n_filters = filters.shape[0]

    # Default groups to n_channels (depthwise)
    if groups is None:
        groups = n_channels

    # Validate groups
    if groups < 1:
        raise ValueError(f"groups must be >= 1, got {groups}")
    if n_channels % groups != 0:
        raise ValueError(
            f"n_channels ({n_channels}) must be divisible by groups ({groups})"
        )
    if n_filters % groups != 0:
        raise ValueError(
            f"n_filters ({n_filters}) must be divisible by groups ({groups})"
        )

    # Check minimum size
    if data.shape[1] < filters.shape[1]:
        raise ValueError(
            f"data must have at least filter_length ({filters.shape[1]}) samples, "
            f"got {data.shape[1]}"
        )

    # Convert to float32 and ensure C-contiguous
    if data.dtype != np.float32:
        data = data.astype(np.float32, order="C")
    elif not data.flags["C_CONTIGUOUS"]:
        data = np.ascontiguousarray(data)

    if filters.dtype != np.float32:
        filters = filters.astype(np.float32, order="C")
    elif not filters.flags["C_CONTIGUOUS"]:
        filters = np.ascontiguousarray(filters)

    # Call SIMD-optimized C extension
    return _correlate_core(data, filters, groups)


__all__ = [
    "upsample",
    "upsample_transposed",
    "correlate",
    "get_simd_info",
    "set_implementation",
]
