import numpy as np
import librosa
from scipy import stats as scipy_stats
import warnings

warnings.filterwarnings('ignore')

# ==========================================
# CONFIGURATION
# ==========================================
TARGET_SR = 48000
TARGET_DURATION = 10.0
TARGET_LENGTH = int(TARGET_SR * TARGET_DURATION)
TRIM_TOP_DB = 20
N_MFCC = 20

# ==========================================
# PREPROCESSING
# ==========================================
def preprocess_audio(file_path):
    y, sr = librosa.load(file_path, sr=TARGET_SR)

    # Trim silence
    y_trimmed, _ = librosa.effects.trim(y, top_db=TRIM_TOP_DB)
    if len(y_trimmed) == 0:
        y_trimmed = y

    # DC offset removal
    y_centered = y_trimmed - np.mean(y_trimmed)

    # Z-score normalization
    std = np.std(y_centered)
    if std > 1e-8:
        y_norm = y_centered / std
    else:
        y_norm = y_centered

    # Fixed length
    if len(y_norm) < TARGET_LENGTH:
        y_norm = np.pad(y_norm, (0, TARGET_LENGTH - len(y_norm)), mode='constant')
    else:
        y_norm = y_norm[:TARGET_LENGTH]

    return y_norm


# ==========================================
# FEATURE EXTRACTION
# ==========================================
def compute_stats(feature_array, prefix):
    result = {}
    if feature_array.ndim == 1:
        feature_array = feature_array.reshape(1, -1)

    for i in range(feature_array.shape[0]):
        row = feature_array[i]
        suffix = f"_{i}" if feature_array.shape[0] > 1 else ""
        result[f"{prefix}{suffix}_mean"] = np.mean(row)
        result[f"{prefix}{suffix}_std"]  = np.std(row)
        result[f"{prefix}{suffix}_min"]  = np.min(row)
        result[f"{prefix}{suffix}_max"]  = np.max(row)
    return result


def extract_ml_features(y, sr=TARGET_SR, n_mfcc=N_MFCC):
    features = {}

    # 1. MFCCs (20 × 4 = 80)
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=n_mfcc)
    features.update(compute_stats(mfccs, "mfcc"))

    # 2. Spectral Centroid (4)
    features.update(compute_stats(librosa.feature.spectral_centroid(y=y, sr=sr), "spectral_centroid"))

    # 3. Spectral Flatness (4)
    features.update(compute_stats(librosa.feature.spectral_flatness(y=y), "spectral_flatness"))

    # 4. Zero-Crossing Rate (4)
    features.update(compute_stats(librosa.feature.zero_crossing_rate(y), "zcr"))

    # 5. RMS Energy (4)
    features.update(compute_stats(librosa.feature.rms(y=y), "rms"))

    # 6. Spectral Rolloff at 85% (4)
    features.update(compute_stats(librosa.feature.spectral_rolloff(y=y, sr=sr, roll_percent=0.85), "spectral_rolloff"))

    # 7. Spectral Bandwidth (4)
    features.update(compute_stats(librosa.feature.spectral_bandwidth(y=y, sr=sr), "spectral_bandwidth"))

    # 8. Spectral Contrast (7 bands × 4 = 28)
    features.update(compute_stats(librosa.feature.spectral_contrast(y=y, sr=sr, n_bands=6), "spectral_contrast"))

    # 9. Chroma STFT (12 × 4 = 48)
    features.update(compute_stats(librosa.feature.chroma_stft(y=y, sr=sr), "chroma"))

    # 10. Spectral Kurtosis (1)
    S = np.abs(librosa.stft(y))
    features["spectral_kurtosis"] = float(scipy_stats.kurtosis(np.mean(S, axis=1)))

    return features