from collections import Counter
from pathlib import Path
import warnings

import librosa
import numpy as np


SAMPLE_RATE = 22050
LABEL_PD = "PD"
LABEL_TPD = "TPD"
VALID_LABELS = {LABEL_PD, LABEL_TPD}
N_MFCC = 13
MIN_DURATION_SECONDS = 0.30
MIN_RMS_FOR_USABLE_AUDIO = 0.001


def get_label_from_filename(file_path):
    """
    Mengambil label dari nama file atau nama folder.

    Urutan pengecekan penting:
    - cek "_tpd" lebih dulu
    - baru cek "_pd"
    """
    path = Path(file_path)
    filename = path.stem.lower()
    parent = path.parent.name.lower()

    if "_tpd" in filename or parent in {"tpd", "not_confident", "tidak_percaya_diri"}:
        return LABEL_TPD
    if "_pd" in filename or parent in {"pd", "confident", "percaya_diri"}:
        return LABEL_PD

    return None


def validate_label(label, file_path):
    if label not in VALID_LABELS:
        raise ValueError(f"Label tidak valid pada {Path(file_path).name}: {label}")


def load_and_preprocess_audio(file_path, sample_rate=SAMPLE_RATE):
    """
    Membaca audio dengan preprocessing konsisten untuk training dan prediksi:
    mono, sample rate 22050 Hz, trim silence, dan normalisasi volume.
    """
    y, sr = librosa.load(file_path, sr=sample_rate, mono=True)

    if y.size == 0:
        raise ValueError(f"Audio kosong: {file_path}")

    y, _ = librosa.effects.trim(y, top_db=30)

    if y.size == 0:
        raise ValueError(f"Audio hanya berisi silence: {file_path}")

    duration = librosa.get_duration(y=y, sr=sr)
    if duration < MIN_DURATION_SECONDS:
        raise ValueError(
            f"Audio terlalu pendek: {duration:.2f} detik. Minimal {MIN_DURATION_SECONDS:.2f} detik."
        )

    rms_value = float(np.sqrt(np.mean(y**2)))
    if rms_value < MIN_RMS_FOR_USABLE_AUDIO:
        raise ValueError(
            f"Audio terlalu pelan/silent. RMS={rms_value:.5f}, "
            f"minimal {MIN_RMS_FOR_USABLE_AUDIO:.5f}."
        )

    max_amplitude = np.max(np.abs(y))
    if max_amplitude > 0:
        y = y / max_amplitude

    return y.astype(np.float32), sr


def mean_std(feature_matrix):
    """
    Mengubah fitur frame-based menjadi statistik tetap.
    Output selalu 1 dimensi dan stabil untuk SVM.
    """
    feature_matrix = np.atleast_2d(feature_matrix)
    return np.concatenate(
        [
            np.mean(feature_matrix, axis=1),
            np.std(feature_matrix, axis=1),
        ]
    )


def extract_pitch_features(y, sr):
    """
    Mengambil ringkasan fundamental frequency (pitch) dengan pyin.
    Jika pitch tidak terdeteksi, nilai pitch dibuat 0 agar fitur tetap konsisten.
    """
    f0, _, _ = librosa.pyin(
        y,
        fmin=librosa.note_to_hz("C2"),
        fmax=librosa.note_to_hz("C7"),
        sr=sr,
    )
    voiced_f0 = f0[~np.isnan(f0)]

    if voiced_f0.size == 0:
        return np.array([0.0, 0.0, 0.0], dtype=np.float32)

    voiced_ratio = voiced_f0.size / f0.size
    return np.array(
        [
            np.mean(voiced_f0),
            np.std(voiced_f0),
            voiced_ratio,
        ],
        dtype=np.float32,
    )


def extract_features(file_path, sample_rate=SAMPLE_RATE):
    """
    Ekstraksi fitur suara yang sama untuk training dan prediksi:
    - MFCC mean dan std
    - RMS Energy mean dan std
    - Zero Crossing Rate mean dan std
    - Spectral Centroid mean dan std
    - Spectral Bandwidth mean dan std
    - Spectral Rolloff mean dan std
    - Pitch/fundamental frequency
    - Durasi suara aktif
    """
    y, sr = load_and_preprocess_audio(file_path, sample_rate=sample_rate)

    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=N_MFCC)
    rms = librosa.feature.rms(y=y)
    zcr = librosa.feature.zero_crossing_rate(y)
    spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=sr)
    spectral_bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sr)
    spectral_rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)
    pitch_features = extract_pitch_features(y, sr)
    active_duration = np.array([librosa.get_duration(y=y, sr=sr)], dtype=np.float32)

    feature_vector = np.concatenate(
        [
            mean_std(mfcc),
            mean_std(rms),
            mean_std(zcr),
            mean_std(spectral_centroid),
            mean_std(spectral_bandwidth),
            mean_std(spectral_rolloff),
            pitch_features,
            active_duration,
        ]
    )

    if feature_vector.ndim != 1:
        raise ValueError("Fitur audio harus 1 dimensi.")
    if not np.all(np.isfinite(feature_vector)):
        raise ValueError("Fitur audio mengandung NaN atau infinity.")

    return feature_vector.astype(np.float32)


def load_dataset(data_dir):
    """
    Membaca semua file .wav pada folder data.
    File tanpa label valid atau file rusak dilewati dengan peringatan.
    """
    data_path = Path(data_dir)
    audio_files = sorted(data_path.rglob("*.wav"))

    if not audio_files:
        raise FileNotFoundError(f"Tidak ada file .wav di folder: {data_path}")

    features = []
    labels = []
    used_files = []

    for audio_file in audio_files:
        label = get_label_from_filename(audio_file)
        if label is None:
            warnings.warn(
                f"File dilewati karena nama/folder tidak mengandung label PD atau TPD: "
                f"{audio_file.name}"
            )
            continue

        try:
            validate_label(label, audio_file)
            features.append(extract_features(audio_file))
            labels.append(label)
            used_files.append(audio_file)
        except Exception as error:
            warnings.warn(f"File dilewati karena gagal diproses: {audio_file.name} ({error})")

    if not features:
        raise ValueError("Tidak ada file audio valid yang berhasil diproses.")

    label_counts = Counter(labels)
    print("\n=== Distribusi Label Dataset ===")
    print(f"PD : {label_counts.get(LABEL_PD, 0)}")
    print(f"TPD: {label_counts.get(LABEL_TPD, 0)}")

    invalid_labels = set(labels) - VALID_LABELS
    if invalid_labels:
        raise ValueError(f"Ditemukan label tidak valid: {sorted(invalid_labels)}")

    return np.array(features), np.array(labels), used_files


def check_dataset_quality(data_dir):
    """
    Mengecek kualitas dataset:
    - jumlah data PD dan TPD
    - durasi setiap audio
    - audio terlalu pendek
    - audio terlalu pelan/silent
    - file rusak
    - rekomendasi file yang perlu direkam ulang
    """
    data_path = Path(data_dir)
    audio_files = sorted(data_path.rglob("*.wav"))
    label_counts = Counter()
    problems = []

    print("\n=== Cek Kualitas Dataset ===")

    for audio_file in audio_files:
        label = get_label_from_filename(audio_file)
        if label is None:
            problems.append((audio_file.name, "Label tidak ditemukan"))
            continue

        label_counts[label] += 1

        try:
            y_raw, sr = librosa.load(audio_file, sr=SAMPLE_RATE, mono=True)
            duration_raw = librosa.get_duration(y=y_raw, sr=sr)
            rms_raw = float(np.sqrt(np.mean(y_raw**2))) if y_raw.size else 0.0

            issue_notes = []
            if duration_raw < MIN_DURATION_SECONDS:
                issue_notes.append(f"terlalu pendek ({duration_raw:.2f} detik)")
            if rms_raw < MIN_RMS_FOR_USABLE_AUDIO:
                issue_notes.append(f"terlalu pelan/silent (RMS={rms_raw:.5f})")

            print(
                f"{audio_file.name} | label={label} | durasi={duration_raw:.2f}s | "
                f"rms={rms_raw:.5f}"
            )

            if issue_notes:
                problems.append((audio_file.name, ", ".join(issue_notes)))
        except Exception as error:
            problems.append((audio_file.name, f"file rusak/gagal dibaca ({error})"))

    print("\nJumlah data:")
    print(f"PD : {label_counts.get(LABEL_PD, 0)}")
    print(f"TPD: {label_counts.get(LABEL_TPD, 0)}")

    print("\nRekomendasi rekam ulang/perbaikan:")
    if not problems:
        print("Tidak ada masalah kualitas audio yang jelas.")
    else:
        for filename, reason in problems:
            print(f"- {filename}: {reason}")

    return problems