import logging

import numpy as np
from scipy.ndimage import uniform_filter1d
from scipy.signal import fftconvolve, firwin2, welch

logger = logging.getLogger(__name__)


def compute_psd(audio: np.ndarray, sr: int) -> tuple[np.ndarray, np.ndarray]:
    """Compute power spectral density using Welch's method.

    For stereo audio, averages PSDs across channels.
    Returns (frequencies, psd).
    """
    if audio.ndim == 2:
        # Average PSD across channels
        freqs, psd0 = welch(audio[:, 0], fs=sr, nperseg=8192)
        _, psd1 = welch(audio[:, 1], fs=sr, nperseg=8192)
        psd = (psd0 + psd1) / 2
    else:
        freqs, psd = welch(audio, fs=sr, nperseg=8192)

    return freqs, psd


def compute_transfer_function(
    psd_ref: np.ndarray,
    psd_user: np.ndarray,
    freqs: np.ndarray,
    smoothing_octaves: float,
    max_gain_db: float,
) -> np.ndarray:
    """Compute a smoothed, clamped magnitude transfer function.

    Returns the transfer function as a linear amplitude array.
    """
    # Amplitude ratio (sqrt because PSD is power)
    with np.errstate(divide="ignore", invalid="ignore"):
        h = np.sqrt(psd_ref / (psd_user + 1e-30))

    # Handle NaN/inf from division
    h = np.where(np.isfinite(h), h, 1.0)

    # Convert to dB
    h_db = 20 * np.log10(h + 1e-10)

    # Smooth in log-frequency domain
    if smoothing_octaves > 0:
        h_db = _octave_smooth(h_db, freqs, smoothing_octaves)

    # Clamp
    h_db = np.clip(h_db, -max_gain_db, max_gain_db)

    # Boundary conditions: DC and Nyquist at unity
    h_db[0] = 0.0
    h_db[-1] = 0.0

    # Convert back to linear
    h_linear = 10 ** (h_db / 20)

    return h_linear


def _octave_smooth(h_db: np.ndarray, freqs: np.ndarray, octave_fraction: float) -> np.ndarray:
    """Apply smoothing in log-frequency domain."""
    if len(freqs) < 2:
        return h_db

    # Work on non-DC frequencies
    log_freqs = np.log2(freqs[1:] + 1e-10)
    bin_spacing = np.median(np.diff(log_freqs))
    if bin_spacing <= 0:
        return h_db

    window_bins = max(1, int(round(octave_fraction / bin_spacing)))

    h_db_smoothed = h_db.copy()
    h_db_smoothed[1:] = uniform_filter1d(h_db[1:], size=window_bins)
    return h_db_smoothed


def apply_eq(
    audio: np.ndarray, sr: int, transfer_function: np.ndarray, freqs: np.ndarray
) -> np.ndarray:
    """Apply EQ via minimum-phase FIR filter convolution.

    Returns the EQ'd audio with the same shape as input.
    """
    # Normalize frequencies to [0, 1] for firwin2
    nyquist = sr / 2
    freqs_norm = freqs / nyquist

    # Ensure frequency endpoints are exactly 0 and 1
    freqs_norm = np.clip(freqs_norm, 0, 1)
    freqs_norm[0] = 0.0
    freqs_norm[-1] = 1.0

    # Design FIR filter
    try:
        fir = firwin2(numtaps=4095, freq=freqs_norm, gain=transfer_function)
    except Exception:
        logger.warning("FIR filter design failed. Skipping EQ for this stem.")
        return audio

    # Apply to each channel
    if audio.ndim == 2:
        result = np.stack(
            [fftconvolve(audio[:, ch], fir, mode="same") for ch in range(audio.shape[1])],
            axis=1,
        )
    else:
        result = fftconvolve(audio, fir, mode="same")

    return result.astype(np.float32)


def match_eq(
    user_stem: np.ndarray,
    ref_stem: np.ndarray,
    sr: int,
    smoothing_octaves: float,
    max_gain_db: float,
) -> np.ndarray:
    """High-level EQ matching: compute PSDs, derive transfer function, apply filter."""
    logger.debug("Computing PSDs...")
    freqs, psd_user = compute_psd(user_stem, sr)
    _, psd_ref = compute_psd(ref_stem, sr)

    # Warn on problematic PSDs
    if np.all(psd_user < 1e-20) or np.all(psd_ref < 1e-20):
        logger.warning("PSD is near-zero. Setting to unity gain.")
        return user_stem

    logger.debug("Computing transfer function (smoothing=%.2f oct, max_gain=%.1f dB)...", smoothing_octaves, max_gain_db)
    tf = compute_transfer_function(psd_ref, psd_user, freqs, smoothing_octaves, max_gain_db)

    logger.debug("Applying EQ filter...")
    return apply_eq(user_stem, sr, tf, freqs)
