from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
import os

import numpy as np
import pandas as pd
import soundfile as sf
from scipy.signal import resample_poly
from math import gcd
from tqdm import tqdm


# ============================================================
# Paths
# ============================================================

AUDIO_ROOT = Path("/home/debarpanb1/TREA_2.0/UrbanSound8K/audio")
CSV_PATH = Path("/home/debarpanb1/TREA_2.0/UrbanSound8K/metadata/UrbanSound8K.csv")

OUTPUT_ROOT = Path("/home/debarpanb1/TREA_2.0/UrbanSound8K_preprocessed_fast")
OUTPUT_AUDIO_ROOT = OUTPUT_ROOT / "audio"
OUTPUT_METADATA_PATH = OUTPUT_ROOT / "UrbanSound8K_with_audio_metadata.csv"

OUTPUT_AUDIO_ROOT.mkdir(parents=True, exist_ok=True)


# ============================================================
# Config
# ============================================================

TARGET_SR = 44100
TARGET_SUBTYPE = "PCM_16"
NORMALIZE_PEAK = True

# Use fewer workers if storage is slow.
NUM_WORKERS = min(16, os.cpu_count() or 4)


# ============================================================
# Helpers
# ============================================================

def safe_class_name(x):
    return str(x).replace(" ", "_").replace("/", "_")


def make_output_filename(row):
    fold = int(row["fold"])
    class_name = safe_class_name(row["class"])
    stem = Path(row["slice_file_name"]).stem
    return f"fold{fold}_{class_name}_{stem}.wav"


def to_mono(audio):
    if audio.ndim == 1:
        return audio
    return audio.mean(axis=1)


def resample_audio(audio, orig_sr, target_sr):
    if orig_sr == target_sr:
        return audio

    factor = gcd(orig_sr, target_sr)
    up = target_sr // factor
    down = orig_sr // factor

    return resample_poly(audio, up, down)


def peak_normalize(audio, eps=1e-8):
    peak = np.max(np.abs(audio)) if len(audio) > 0 else 0.0
    if peak < eps:
        return audio
    return audio / peak * 0.98


def process_one(row_dict):
    fold = int(row_dict["fold"])
    filename = row_dict["slice_file_name"]

    original_audio_path = AUDIO_ROOT / f"fold{fold}" / filename
    output_filename = make_output_filename(row_dict)
    clean_audio_path = OUTPUT_AUDIO_ROOT / f"fold{fold}" / output_filename
    clean_audio_path.parent.mkdir(parents=True, exist_ok=True)

    row_dict["original_audio_path"] = str(original_audio_path)
    row_dict["clean_audio_path"] = str(clean_audio_path)

    row_dict["preprocess_target_sample_rate"] = TARGET_SR
    row_dict["preprocess_target_channels"] = 1
    row_dict["preprocess_target_subtype"] = TARGET_SUBTYPE
    row_dict["preprocess_peak_normalize"] = NORMALIZE_PEAK
    row_dict["preprocess_success"] = False
    row_dict["preprocess_error"] = ""

    if not original_audio_path.exists():
        row_dict["preprocess_error"] = "missing_audio_file"
        return row_dict

    try:
        # Original metadata
        info = sf.info(str(original_audio_path))

        row_dict["original_sample_rate"] = info.samplerate
        row_dict["original_channels"] = info.channels
        row_dict["original_frames"] = info.frames
        row_dict["original_duration"] = info.frames / info.samplerate
        row_dict["original_format"] = info.format
        row_dict["original_subtype"] = info.subtype

        row_dict["csv_slice_duration"] = float(row_dict["end"]) - float(row_dict["start"])

        # Fast read
        audio, sr = sf.read(str(original_audio_path), dtype="float32", always_2d=False)

        # Mono
        audio = to_mono(audio)

        # Resample only if needed
        audio = resample_audio(audio, sr, TARGET_SR)

        # Normalize
        if NORMALIZE_PEAK:
            audio = peak_normalize(audio)

        # Save
        sf.write(str(clean_audio_path), audio, TARGET_SR, subtype=TARGET_SUBTYPE)

        # Clean metadata
        clean_info = sf.info(str(clean_audio_path))

        row_dict["clean_sample_rate"] = clean_info.samplerate
        row_dict["clean_channels"] = clean_info.channels
        row_dict["clean_frames"] = clean_info.frames
        row_dict["clean_duration"] = clean_info.frames / clean_info.samplerate
        row_dict["clean_format"] = clean_info.format
        row_dict["clean_subtype"] = clean_info.subtype

        row_dict["preprocess_action"] = "soundfile_scipy_resample_to_44100_mono_pcm16_peaknorm"
        row_dict["preprocess_success"] = True

    except Exception as e:
        row_dict["preprocess_error"] = repr(e)

    return row_dict


def main():
    df = pd.read_csv(CSV_PATH)

    records = df.to_dict("records")
    results = []

    print(f"Processing {len(records)} files with {NUM_WORKERS} workers...")

    with ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
        futures = [executor.submit(process_one, row) for row in records]

        for future in tqdm(as_completed(futures), total=len(futures)):
            results.append(future.result())

    out_df = pd.DataFrame(results)
    out_df.to_csv(OUTPUT_METADATA_PATH, index=False)

    print("\nDone.")
    print(f"Preprocessed audio saved to: {OUTPUT_AUDIO_ROOT}")
    print(f"Metadata saved to: {OUTPUT_METADATA_PATH}")

    print("\nSuccess count:")
    print(out_df["preprocess_success"].value_counts(dropna=False))

    print("\nOriginal sample rates:")
    print(out_df["original_sample_rate"].value_counts(dropna=False).sort_index())

    print("\nOriginal channels:")
    print(out_df["original_channels"].value_counts(dropna=False).sort_index())

    print("\nOriginal subtypes:")
    print(out_df["original_subtype"].value_counts(dropna=False))


if __name__ == "__main__":
    main()