"""
Spectral feature calculation functions for audio conditioning.

All functions follow the same pattern:
1. Convert to mono if stereo
2. Window audio at target_rate (configurable, typically 25Hz)
3. Calculate feature for each window
4. Apply smoothing with larger kernel for noise reduction
5. Return feature sequence at target_rate resolution

The target_rate typically matches semantic token rate for 1:1 correspondence.
"""

import numpy as np
import random


def warp_sequence(sequence, warp_ratio=0.5, num_anchor_points=None):
    """Apply time-warping to any 1D sequence for data augmentation.

    This is a pure time-warping function that changes temporal progression
    while preserving the values from the original sequence. The contour shape
    changes but values are resampled from the original.

    Args:
        sequence: Input 1D numpy array (any feature sequence)
        warp_ratio: Warp strength [0,1]. 0=no warp, 1=full warp to target positions
        num_anchor_points: Number of warp control points (default: random 3-20)

    Returns:
        np.ndarray: Time-warped sequence (same length, values from original)

    Algorithm:
        1. Sample N anchor indices (sorted, including 0 and len-1)
        2. Sample N target indices (sorted, including 0 and len-1)
        3. Interpolate actual warp: warped[i] = anchor[i] + α*(target[i] - anchor[i])
        4. For each output position, find where it maps in the warped timeline
        5. Resample from original sequence at that position

    Example with warp_ratio=1.0:
        anchor = [0, 50, 100]  (original control points)
        target = [0, 30, 100]  (desired control points - compresses middle)

        Output position 40:
        - In warped space: falls between warp[1]=30 and warp[2]=100
        - Proportion: (40-30)/(100-30) = 0.143
        - Maps back to original: 50 + 0.143*(100-50) = 57.15
        - Sample from sequence[57.15] via interpolation
        - Result: temporal compression - middle of sequence plays faster
    """
    if sequence is None or len(sequence) == 0:
        return sequence

    if warp_ratio <= 0:
        return sequence.copy()

    seq_len = len(sequence)
    if seq_len < 3:
        return sequence.copy()

    # Determine number of anchor points
    if num_anchor_points is None:
        num_anchor_points = random.randint(3, min(20, seq_len))

    # Sample anchor points (original control point positions)
    if num_anchor_points == 2:
        anchor_indices = np.array([0, seq_len - 1], dtype=float)
    else:
        interior_count = num_anchor_points - 2
        interior_points = np.sort(
            np.random.choice(range(1, seq_len - 1), size=interior_count, replace=False)
        )
        anchor_indices = np.concatenate([[0], interior_points, [seq_len - 1]]).astype(float)

    # Sample target points (desired control point positions after warp)
    if num_anchor_points == 2:
        target_indices = np.array([0, seq_len - 1], dtype=float)
    else:
        interior_count = num_anchor_points - 2
        interior_points = np.sort(
            np.random.choice(range(1, seq_len - 1), size=interior_count, replace=False)
        )
        target_indices = np.concatenate([[0], interior_points, [seq_len - 1]]).astype(float)

    # Interpolate warped control points based on warp_ratio
    # 0 = no warp (stay at anchor), 1 = full warp (move to target)
    warped_indices = anchor_indices + warp_ratio * (target_indices - anchor_indices)

    # Create output sequence by inverse time warping
    output_sequence = np.zeros_like(sequence)

    for i in range(seq_len):
        # Find which segment this OUTPUT position falls into in WARPED space
        j = np.searchsorted(warped_indices, i, side="right") - 1
        j = min(max(j, 0), len(warped_indices) - 2)

        # Get warped segment bounds
        warp_start = warped_indices[j]
        warp_end = warped_indices[j + 1]
        anchor_start = anchor_indices[j]
        anchor_end = anchor_indices[j + 1]

        # Find position in warped segment
        if warp_end > warp_start:
            t = (i - warp_start) / (warp_end - warp_start)
        else:
            t = 0.0

        # Map back to original sequence position
        orig_pos = anchor_start + t * (anchor_end - anchor_start)
        orig_pos = np.clip(orig_pos, 0, seq_len - 1)

        # Sample from original sequence using linear interpolation
        pos_floor = int(np.floor(orig_pos))
        pos_ceil = int(np.ceil(orig_pos))

        if pos_floor == pos_ceil:
            output_sequence[i] = sequence[pos_floor]
        else:
            alpha = orig_pos - pos_floor
            output_sequence[i] = (1 - alpha) * sequence[pos_floor] + alpha * sequence[pos_ceil]

    return output_sequence


def calculate_loudness_seq(
    audio_data,
    sample_rate=24000,
    target_rate=25,
    smoothing_kernel_size=7,
    normalize=False,
    apply_warp=False,
    warp_ratio=0.5,
):
    """Calculate loudness sequence from audio data at specified rate.

    Args:
        audio_data: Audio data array (can be mono or stereo)
        sample_rate: Original sample rate of audio data
        target_rate: Target resolution in Hz (default 25, typically matches semantic token rate)
        smoothing_kernel_size: Size of moving average kernel (default 7 for noise reduction)
        normalize: If True, normalize to [0, 1] range using min-max normalization
        apply_warp: If True, apply time-warping BEFORE smoothing
        warp_ratio: Time-warp strength [0,1], only used if apply_warp=True

    Returns:
        np.ndarray: Loudness values at target_rate resolution (normalized if requested)
    """
    if audio_data.ndim > 1:
        # Convert to mono if stereo/multi-channel
        audio_data = np.mean(audio_data, axis=0)

    # Calculate window size for target rate
    window_size = int(sample_rate / target_rate)

    # Pad to ensure even windows
    num_windows = int(np.ceil(len(audio_data) / window_size))
    padded_length = num_windows * window_size
    padded_audio = np.pad(audio_data, (0, padded_length - len(audio_data)), mode="constant")

    # Reshape into windows and calculate RMS (root mean square) loudness
    windowed = padded_audio.reshape(num_windows, window_size)
    rms_loudness = np.sqrt(np.mean(windowed**2, axis=1))

    # Apply time-warping BEFORE smoothing (if enabled)
    if apply_warp and warp_ratio > 0:
        rms_loudness = warp_sequence(rms_loudness, warp_ratio)

    # Apply smoothing with a moving average using numpy
    kernel = np.ones(smoothing_kernel_size) / smoothing_kernel_size
    smoothed_loudness = np.convolve(rms_loudness, kernel, mode="same")

    # Normalize to [0, 1] if requested
    if normalize:
        min_val = np.min(smoothed_loudness)
        max_val = np.max(smoothed_loudness)
        if max_val > min_val:
            smoothed_loudness = (smoothed_loudness - min_val) / (max_val - min_val)
        else:
            # All values are the same, set to 0
            smoothed_loudness = np.zeros_like(smoothed_loudness)

    return smoothed_loudness


def calculate_spectral_centroid_seq(
    audio_data,
    sample_rate=24000,
    target_rate=25,
    smoothing_kernel_size=7,
    normalize=False,
    loudness_seq=None,
    apply_warp=False,
    warp_ratio=0.5,
):
    """Calculate spectral centroid sequence from audio data at specified rate.

    The spectral centroid represents the "center of mass" of the spectrum,
    indicating the brightness of the sound. Higher values = brighter sound.

    Args:
        audio_data: Audio data array (can be mono or stereo)
        sample_rate: Original sample rate of audio data
        target_rate: Target resolution in Hz (default 25, typically matches semantic token rate)
        smoothing_kernel_size: Size of moving average kernel (default 7 for noise reduction)
        normalize: If True, normalize to [0, 1] range using min-max normalization
        loudness_seq: Optional loudness sequence for silence detection during normalization.
                     Must have same length as output. If provided with normalize=True,
                     uses loudness > 0.2 threshold to detect active frames.
        apply_warp: If True, apply time-warping BEFORE smoothing
        warp_ratio: Time-warp strength [0,1], only used if apply_warp=True

    Returns:
        np.ndarray: Spectral centroid values in Hz at target_rate resolution (normalized if requested)
    """
    if audio_data.ndim > 1:
        # Convert to mono if stereo/multi-channel
        audio_data = np.mean(audio_data, axis=0)

    # Calculate window size for target rate
    window_size = int(sample_rate / target_rate)

    # Pad to ensure even windows
    num_windows = int(np.ceil(len(audio_data) / window_size))
    padded_length = num_windows * window_size
    padded_audio = np.pad(audio_data, (0, padded_length - len(audio_data)), mode="constant")

    # Reshape into windows
    windowed = padded_audio.reshape(num_windows, window_size)

    # Apply Hamming window to reduce spectral leakage
    hamming_window = np.hamming(window_size)

    centroids = []
    for i in range(num_windows):
        # Apply window function
        windowed_frame = windowed[i] * hamming_window

        # Compute FFT magnitude spectrum
        spectrum = np.abs(np.fft.rfft(windowed_frame))

        # Frequency bins in Hz
        freqs = np.fft.rfftfreq(window_size, 1 / sample_rate)

        # Handle edge case: all zeros (silence)
        if np.sum(spectrum) < 1e-10:
            centroids.append(0.0)
            continue

        # Calculate spectral centroid: weighted mean of frequencies
        centroid = np.sum(freqs * spectrum) / np.sum(spectrum)
        centroids.append(centroid)

    centroids = np.array(centroids)

    # Replace any NaN or inf values with 0
    centroids = np.nan_to_num(centroids, nan=0.0, posinf=0.0, neginf=0.0)

    # Apply time-warping BEFORE smoothing (if enabled)
    if apply_warp and warp_ratio > 0:
        centroids = warp_sequence(centroids, warp_ratio)

    # Apply smoothing with a moving average
    kernel = np.ones(smoothing_kernel_size) / smoothing_kernel_size
    smoothed_centroids = np.convolve(centroids, kernel, mode="same")

    # Normalize to [0, 1] if requested
    if normalize:
        if loudness_seq is not None:
            # Use loudness-based silence detection
            if len(loudness_seq) != len(smoothed_centroids):
                raise ValueError(
                    f"loudness_seq length ({len(loudness_seq)}) must match "
                    f"spectral_centroid length ({len(smoothed_centroids)}). "
                    f"Ensure both are calculated at the same target_rate."
                )

            # Active frames are where loudness > 0.2
            active_mask = loudness_seq > 0.2
        else:
            # Fallback: Filter out silence frames (very low centroids, typically < 100 Hz)
            # to avoid skewing normalization with silence
            silence_threshold = 100  # Hz
            active_mask = smoothed_centroids > silence_threshold

        if np.any(active_mask):
            # Calculate min/max only from active (non-silence) frames
            # Also exclude first and last 5% of frames to avoid extreme outliers
            total_frames = len(smoothed_centroids)
            trim_count = int(0.05 * total_frames)

            # Get indices of active frames
            active_indices = np.where(active_mask)[0]

            # Filter out indices in the first and last 5%
            middle_mask = (active_indices >= trim_count) & (active_indices < total_frames - trim_count)
            middle_active_indices = active_indices[middle_mask]

            if len(middle_active_indices) > 0:
                # Use middle active frames for min/max calculation
                min_val = np.min(smoothed_centroids[middle_active_indices])
                max_val = np.max(smoothed_centroids[middle_active_indices])
            else:
                # Fall back to all active frames if middle is empty
                min_val = np.min(smoothed_centroids[active_mask])
                max_val = np.max(smoothed_centroids[active_mask])

            if max_val > min_val:
                smoothed_centroids = (smoothed_centroids - min_val) / (max_val - min_val)
                # Clip to [0, 1] in case silence frames go below 0
                smoothed_centroids = np.clip(smoothed_centroids, 0.0, 1.0)
            else:
                # All active values are the same, set to 0.5
                smoothed_centroids = np.full_like(smoothed_centroids, 0.5)
        else:
            # All silence, set to 0
            smoothed_centroids = np.zeros_like(smoothed_centroids)

    return smoothed_centroids


def calculate_spectral_complexity_seq(
    audio_data,
    sample_rate=24000,
    target_rate=25,
    smoothing_kernel_size=7,
    normalize=False,
    loudness_seq=None,
    apply_warp=False,
    warp_ratio=0.5,
):
    """Calculate spectral complexity (entropy) sequence from audio data at specified rate.

    Spectral complexity measures the disorder or richness of the frequency content.
    Higher values indicate more complex/rich harmonic structure.
    Returns values in 0-1 range where:
    - 0 = pure tone (simple)
    - 1 = white noise (maximally complex)

    Args:
        audio_data: Audio data array (can be mono or stereo)
        sample_rate: Original sample rate of audio data
        target_rate: Target resolution in Hz (default 25, typically matches semantic token rate)
        smoothing_kernel_size: Size of moving average kernel (default 7 for noise reduction)
        normalize: If True, normalize to [0, 1] range using min-max normalization
        loudness_seq: Optional loudness sequence for silence detection during normalization.
                     Must have same length as output. If provided with normalize=True,
                     uses loudness > 0.2 threshold to detect active frames.
        apply_warp: If True, apply time-warping BEFORE smoothing
        warp_ratio: Time-warp strength [0,1], only used if apply_warp=True

    Returns:
        np.ndarray: Spectral complexity values (0-1) at target_rate resolution (normalized if requested)
    """
    if audio_data.ndim > 1:
        # Convert to mono if stereo/multi-channel
        audio_data = np.mean(audio_data, axis=0)

    # Calculate window size for target rate
    window_size = int(sample_rate / target_rate)

    # Pad to ensure even windows
    num_windows = int(np.ceil(len(audio_data) / window_size))
    padded_length = num_windows * window_size
    padded_audio = np.pad(audio_data, (0, padded_length - len(audio_data)), mode="constant")

    # Reshape into windows
    windowed = padded_audio.reshape(num_windows, window_size)

    # Apply Hamming window to reduce spectral leakage
    hamming_window = np.hamming(window_size)

    complexities = []
    for i in range(num_windows):
        # Apply window function
        windowed_frame = windowed[i] * hamming_window

        # Compute power spectrum
        spectrum = np.abs(np.fft.rfft(windowed_frame)) ** 2

        # Handle edge case: all zeros (silence)
        total_power = np.sum(spectrum)
        if total_power < 1e-10:
            complexities.append(0.0)
            continue

        # Normalize to probability distribution
        p = spectrum / total_power

        # Add small epsilon to avoid log(0)
        epsilon = 1e-10
        p = p + epsilon
        p = p / np.sum(p)  # Renormalize after adding epsilon

        # Calculate spectral entropy
        entropy = -np.sum(p * np.log2(p))

        # Normalize by maximum possible entropy (log2 of number of bins)
        # This gives us a 0-1 range
        max_entropy = np.log2(len(p))
        normalized_entropy = entropy / max_entropy if max_entropy > 0 else 0.0

        complexities.append(normalized_entropy)

    complexities = np.array(complexities)

    # Replace any NaN or inf values with 0
    complexities = np.nan_to_num(complexities, nan=0.0, posinf=0.0, neginf=0.0)

    # Apply time-warping BEFORE smoothing (if enabled)
    if apply_warp and warp_ratio > 0:
        complexities = warp_sequence(complexities, warp_ratio)

    # Apply smoothing with a moving average
    kernel = np.ones(smoothing_kernel_size) / smoothing_kernel_size
    smoothed_complexities = np.convolve(complexities, kernel, mode="same")

    # Clip to [0, 1] range (in case smoothing causes slight overflow)
    smoothed_complexities = np.clip(smoothed_complexities, 0.0, 1.0)

    # Normalize to [0, 1] if requested (stretches the range to use full [0,1])
    if normalize:
        if loudness_seq is not None:
            # Use loudness-based silence detection
            if len(loudness_seq) != len(smoothed_complexities):
                raise ValueError(
                    f"loudness_seq length ({len(loudness_seq)}) must match "
                    f"spectral_complexity length ({len(smoothed_complexities)}). "
                    f"Ensure both are calculated at the same target_rate."
                )

            # Active frames are where loudness > 0.2
            active_mask = loudness_seq > 0.2
        else:
            # Fallback: Filter out silence frames (very low complexity, typically < 0.1)
            # to avoid skewing normalization with silence
            silence_threshold = 0.1
            active_mask = smoothed_complexities > silence_threshold

        if np.any(active_mask):
            # Calculate min/max only from active (non-silence) frames
            # Also exclude first and last 5% of frames to avoid extreme outliers
            total_frames = len(smoothed_complexities)
            trim_count = int(0.05 * total_frames)

            # Get indices of active frames
            active_indices = np.where(active_mask)[0]

            # Filter out indices in the first and last 5%
            middle_mask = (active_indices >= trim_count) & (active_indices < total_frames - trim_count)
            middle_active_indices = active_indices[middle_mask]

            if len(middle_active_indices) > 0:
                # Use middle active frames for min/max calculation
                min_val = np.min(smoothed_complexities[middle_active_indices])
                max_val = np.max(smoothed_complexities[middle_active_indices])
            else:
                # Fall back to all active frames if middle is empty
                min_val = np.min(smoothed_complexities[active_mask])
                max_val = np.max(smoothed_complexities[active_mask])

            if max_val > min_val:
                smoothed_complexities = (smoothed_complexities - min_val) / (max_val - min_val)
                # Clip to [0, 1] in case silence frames go below 0
                smoothed_complexities = np.clip(smoothed_complexities, 0.0, 1.0)
            else:
                # All active values are the same, set to 0.5
                smoothed_complexities = np.full_like(smoothed_complexities, 0.5)
        else:
            # All silence, set to 0
            smoothed_complexities = np.zeros_like(smoothed_complexities)

    return smoothed_complexities