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 = 2.00
MIN_RMS_FOR_USABLE_AUDIO = 0.001
QUIET_RMS_THRESHOLD = 0.003
CLIPPING_AMPLITUDE_THRESHOLD = 0.99
CLIPPING_RATIO_THRESHOLD = 0.01
FRAME_LENGTH = 2048
HOP_LENGTH = 512
FEATURE_NAMES = (
    [f"mfcc_{index}_mean" for index in range(1, N_MFCC + 1)]
    + [f"mfcc_{index}_std" for index in range(1, N_MFCC + 1)]
    + [f"delta_mfcc_{index}_mean" for index in range(1, N_MFCC + 1)]
    + [f"delta_mfcc_{index}_std" for index in range(1, N_MFCC + 1)]
    + ["rms_normalized_mean", "rms_normalized_std"]
    + ["rms_mean", "rms_std", "peak_amplitude", "clipping_ratio", "energy_stability"]
    + ["pitch_mean", "pitch_std", "pitch_range", "pitch_stability", "pitch_variation"]
    + ["zcr_mean", "zcr_std"]
    + ["spectral_centroid_mean", "spectral_centroid_std"]
    + ["spectral_bandwidth_mean", "spectral_bandwidth_std"]
    + ["spectral_rolloff_mean", "spectral_rolloff_std"]
    + [
        "active_duration",
        "silence_duration",
        "silence_ratio",
        "number_of_pauses",
        "average_pause_duration",
        "speech_activity_ratio",
    ]
)


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_audio_raw(file_path, sample_rate=SAMPLE_RATE):
    y, sr = librosa.load(file_path, sr=sample_rate, mono=True)

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

    return y.astype(np.float32), sr


def calculate_audio_quality(y, sr):
    duration = float(librosa.get_duration(y=y, sr=sr)) if y.size else 0.0
    rms_frames = librosa.feature.rms(y=y, frame_length=FRAME_LENGTH, hop_length=HOP_LENGTH)[0] if y.size else np.array([0.0])
    rms_mean = float(np.mean(rms_frames))
    rms_std = float(np.std(rms_frames))
    peak_amplitude = float(np.max(np.abs(y))) if y.size else 0.0
    clipping_ratio = float(np.mean(np.abs(y) >= CLIPPING_AMPLITUDE_THRESHOLD)) if y.size else 0.0

    return {
        "duration": duration,
        "rms_mean": rms_mean,
        "rms_std": rms_std,
        "peak_amplitude": peak_amplitude,
        "clipping_ratio": clipping_ratio,
        "is_clipped": bool(clipping_ratio > CLIPPING_RATIO_THRESHOLD or peak_amplitude >= CLIPPING_AMPLITUDE_THRESHOLD),
        "is_too_quiet": bool(rms_mean < QUIET_RMS_THRESHOLD),
        "is_too_short": bool(duration < MIN_DURATION_SECONDS),
    }


def validate_audio_for_prediction(audio_quality):
    if audio_quality["is_too_short"]:
        raise ValueError("Audio terlalu pendek, silakan rekam ulang.")
    if audio_quality["is_clipped"]:
        raise ValueError("Audio terlalu keras/pecah, silakan rekam ulang dengan volume normal.")


def normalize_audio(y, target_peak=0.95):
    peak_amplitude = np.max(np.abs(y)) if y.size else 0.0
    if peak_amplitude <= 0:
        return y.astype(np.float32)
    return (y / peak_amplitude * target_peak).astype(np.float32)


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

    if validate_quality:
        validate_audio_for_prediction(audio_quality)

    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 validate_quality and 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}."
        )

    y = normalize_audio(y)

    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 {
            "pitch_mean": 0.0,
            "pitch_std": 0.0,
            "pitch_range": 0.0,
            "pitch_stability": 0.0,
            "pitch_variation": 0.0,
        }

    pitch_mean = float(np.mean(voiced_f0))
    pitch_std = float(np.std(voiced_f0))
    pitch_range = float(np.max(voiced_f0) - np.min(voiced_f0))
    pitch_variation = float(pitch_std / (pitch_mean + 1e-8))
    pitch_stability = float(1.0 / (1.0 + pitch_variation))

    return {
        "pitch_mean": pitch_mean,
        "pitch_std": pitch_std,
        "pitch_range": pitch_range,
        "pitch_stability": pitch_stability,
        "pitch_variation": pitch_variation,
    }


def extract_pause_features(y, sr):
    y_trimmed, _ = librosa.effects.trim(y, top_db=30)
    if y_trimmed.size == 0:
        y_trimmed = y

    total_duration = float(librosa.get_duration(y=y_trimmed, sr=sr))
    rms_frames = librosa.feature.rms(y=y_trimmed, frame_length=FRAME_LENGTH, hop_length=HOP_LENGTH)[0]

    if rms_frames.size == 0 or total_duration <= 0:
        return {
            "active_duration": 0.0,
            "silence_duration": total_duration,
            "silence_ratio": 1.0,
            "number_of_pauses": 0.0,
            "average_pause_duration": 0.0,
            "speech_activity_ratio": 0.0,
        }

    max_rms = float(np.max(rms_frames))
    median_rms = float(np.median(rms_frames))
    adaptive_threshold = max(0.0006, max_rms * 0.06, median_rms * 0.55)
    active_frames = rms_frames > adaptive_threshold
    frame_duration = HOP_LENGTH / sr

    # Lubang diam yang sangat pendek masih dianggap bagian dari artikulasi normal.
    min_gap_frames = int(np.ceil(0.18 / frame_duration))
    smoothed_active_frames = active_frames.copy()
    inactive_run = 0
    inactive_start = 0
    for index, is_active in enumerate(active_frames):
        if not is_active:
            if inactive_run == 0:
                inactive_start = index
            inactive_run += 1
        else:
            if 0 < inactive_run < min_gap_frames:
                smoothed_active_frames[inactive_start:index] = True
            inactive_run = 0

    if 0 < inactive_run < min_gap_frames:
        smoothed_active_frames[inactive_start:] = True

    active_duration = float(min(total_duration, np.sum(smoothed_active_frames) * frame_duration))
    silence_duration = float(max(0.0, total_duration - active_duration))
    silence_ratio = float(silence_duration / (total_duration + 1e-8))
    speech_activity_ratio = float(active_duration / (total_duration + 1e-8))

    pauses = []
    inactive_run = 0
    for is_active in smoothed_active_frames:
        if is_active:
            if inactive_run * frame_duration >= 0.30:
                pauses.append(inactive_run * frame_duration)
            inactive_run = 0
        else:
            inactive_run += 1

    if inactive_run * frame_duration >= 0.30:
        pauses.append(inactive_run * frame_duration)

    return {
        "active_duration": active_duration,
        "silence_duration": silence_duration,
        "silence_ratio": silence_ratio,
        "number_of_pauses": float(len(pauses)),
        "average_pause_duration": float(np.mean(pauses)) if pauses else 0.0,
        "speech_activity_ratio": speech_activity_ratio,
    }


def calculate_voice_indicator_scores(raw_quality, pitch_features, pause_features):
    rms_mean = raw_quality["rms_mean"]
    rms_std = raw_quality["rms_std"]
    clipping_ratio = raw_quality["clipping_ratio"]
    peak_amplitude = raw_quality["peak_amplitude"]

    volume_strength = np.clip(
        np.log10((rms_mean + 1e-8) / QUIET_RMS_THRESHOLD + 1.0) / np.log10(14.0),
        0.0,
        1.0,
    )
    energy_stability = 1.0 / (1.0 + (rms_std / (rms_mean + 1e-8)))
    clipping_penalty = np.clip((clipping_ratio / CLIPPING_RATIO_THRESHOLD) + max(0.0, peak_amplitude - 0.95) * 10.0, 0.0, 1.0)
    volume_score = float(np.clip((0.65 * volume_strength + 0.35 * energy_stability) * (1.0 - clipping_penalty), 0.0, 1.0))

    pitch_variation = pitch_features["pitch_variation"]
    pitch_range = pitch_features["pitch_range"]
    natural_variation = np.clip(pitch_range / 180.0, 0.0, 1.0)
    not_flat = np.clip(pitch_variation / 0.08, 0.0, 1.0)
    not_shaky = 1.0 - np.clip(max(0.0, pitch_variation - 0.35) / 0.45, 0.0, 1.0)
    intonation_score = float(np.clip(0.35 * natural_variation + 0.30 * not_flat + 0.35 * not_shaky, 0.0, 1.0))

    silence_ratio = pause_features["silence_ratio"]
    average_pause_duration = pause_features["average_pause_duration"]
    number_of_pauses = pause_features["number_of_pauses"]
    silence_penalty = np.clip((silence_ratio - 0.12) / 0.38, 0.0, 1.0)
    pause_count_penalty = np.clip(number_of_pauses / 5.0, 0.0, 1.0)
    pause_duration_penalty = np.clip(average_pause_duration / 0.90, 0.0, 1.0)
    pause_score = 1.0 - (0.55 * silence_penalty + 0.25 * pause_count_penalty + 0.20 * pause_duration_penalty)

    return {
        "volume_score": float(np.clip(volume_score, 0.0, 1.0)),
        "intonation_score": float(np.clip(intonation_score, 0.0, 1.0)),
        "pause_score": float(np.clip(pause_score, 0.0, 1.0)),
        "speech_activity_ratio": float(np.clip(pause_features["speech_activity_ratio"], 0.0, 1.0)),
        "silence_ratio": float(np.clip(silence_ratio, 0.0, 1.0)),
    }


def build_prediction_explanation(predicted_label, confidence, indicators):
    weak_points = []
    strong_points = []

    if indicators["volume_score"] >= 0.65:
        strong_points.append("volume stabil")
    else:
        weak_points.append("volume kurang stabil atau kurang jelas")

    if indicators["intonation_score"] >= 0.65:
        strong_points.append("intonasi cukup bervariasi")
    else:
        weak_points.append("intonasi kurang stabil")

    if indicators["pause_score"] >= 0.65:
        strong_points.append("jeda bicara sedikit dan wajar")
    else:
        weak_points.append("jeda bicara cukup banyak")

    if confidence < 0.60:
        reason = " dan ".join(weak_points[:2]) if weak_points else "indikator suara saling berdekatan"
        return f"Model belum cukup yakin karena {reason}."

    if predicted_label == LABEL_PD:
        reason = ", ".join(strong_points[:2])
        if len(strong_points) > 2:
            reason += f", dan {strong_points[2]}"
        return f"Suara terdeteksi percaya diri karena {reason}."

    reason = " dan ".join(weak_points[:2]) if weak_points else "kombinasi indikator belum menunjukkan kestabilan yang cukup"
    return f"Suara terdeteksi tidak percaya diri karena {reason}."


def analyze_audio(file_path, sample_rate=SAMPLE_RATE, validate_quality=False):
    raw_y, sr = load_audio_raw(file_path, sample_rate=sample_rate)
    raw_quality = calculate_audio_quality(raw_y, sr)

    if validate_quality:
        validate_audio_for_prediction(raw_quality)

    y, sr = load_and_preprocess_audio(file_path, sample_rate=sample_rate, validate_quality=False)

    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=N_MFCC, hop_length=HOP_LENGTH)
    delta_mfcc = librosa.feature.delta(mfcc)
    rms_normalized = librosa.feature.rms(y=y, frame_length=FRAME_LENGTH, hop_length=HOP_LENGTH)
    rms_raw = librosa.feature.rms(y=raw_y, frame_length=FRAME_LENGTH, hop_length=HOP_LENGTH)[0]
    zcr = librosa.feature.zero_crossing_rate(y, frame_length=FRAME_LENGTH, hop_length=HOP_LENGTH)
    spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=sr, hop_length=HOP_LENGTH)
    spectral_bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sr, hop_length=HOP_LENGTH)
    spectral_rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr, hop_length=HOP_LENGTH)
    pitch_features = extract_pitch_features(y, sr)
    pause_features = extract_pause_features(raw_y, sr)

    rms_mean = float(np.mean(rms_raw))
    rms_std = float(np.std(rms_raw))
    energy_stability = float(1.0 / (1.0 + (rms_std / (rms_mean + 1e-8))))

    scalar_features = {
        "rms_mean": rms_mean,
        "rms_std": rms_std,
        "peak_amplitude": raw_quality["peak_amplitude"],
        "clipping_ratio": raw_quality["clipping_ratio"],
        "energy_stability": energy_stability,
        **pitch_features,
        **pause_features,
    }

    feature_vector = np.concatenate(
        [
            mean_std(mfcc),
            mean_std(delta_mfcc),
            mean_std(rms_normalized),
            np.array(
                [
                    scalar_features["rms_mean"],
                    scalar_features["rms_std"],
                    scalar_features["peak_amplitude"],
                    scalar_features["clipping_ratio"],
                    scalar_features["energy_stability"],
                    scalar_features["pitch_mean"],
                    scalar_features["pitch_std"],
                    scalar_features["pitch_range"],
                    scalar_features["pitch_stability"],
                    scalar_features["pitch_variation"],
                ],
                dtype=np.float32,
            ),
            mean_std(zcr),
            mean_std(spectral_centroid),
            mean_std(spectral_bandwidth),
            mean_std(spectral_rolloff),
            np.array(
                [
                    scalar_features["active_duration"],
                    scalar_features["silence_duration"],
                    scalar_features["silence_ratio"],
                    scalar_features["number_of_pauses"],
                    scalar_features["average_pause_duration"],
                    scalar_features["speech_activity_ratio"],
                ],
                dtype=np.float32,
            ),
        ]
    )

    if feature_vector.ndim != 1:
        raise ValueError("Fitur audio harus 1 dimensi.")
    if feature_vector.size != len(FEATURE_NAMES):
        raise ValueError(
            f"Jumlah fitur tidak konsisten: {feature_vector.size}, seharusnya {len(FEATURE_NAMES)}."
        )
    if not np.all(np.isfinite(feature_vector)):
        raise ValueError("Fitur audio mengandung NaN atau infinity.")

    indicators = calculate_voice_indicator_scores(raw_quality, pitch_features, pause_features)

    return {
        "features": feature_vector.astype(np.float32),
        "audio_quality": raw_quality,
        "voice_indicators": indicators,
        "feature_details": scalar_features,
    }


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
    """
    return analyze_audio(file_path, sample_rate=sample_rate)["features"]


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
