# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.9.9 64-bit (''venv'': venv)'
#     language: python
#     name: python3
# ---

# +
import pandas as pd
import numpy as np

import librosa
import torchaudio
import torchaudio.transforms as T

import matplotlib.pyplot as plt

import os

from concurrent.futures import ThreadPoolExecutor

from pathlib import Path

import dask.dataframe as dd
from dask.multiprocessing import get
# -

csv_path = "../datasets/AnimalSoundFull.csv"
df = pd.read_csv(csv_path)
print("Columns:\n\t" + "\n\t".join(df.columns))

# +
# To generate ids, if splitting occurs
# def uniqueIDGenerator() -> int:
#     cnt = 0
#     while True:
#         yield cnt
#         cnt += 1

# +
sound_p = Path("../data")
spec_p = Path("../spectrograms/full_dataset/images")
current_spectograms = os.listdir(spec_p)


def generateSpectogram(row):
    if f"{row['gbifID']}.jpg" in current_spectograms:
        return
    if "foo" in row.file_name:
        return

    # Waveform, samplerate
    try:
        wav, sr = torchaudio.load(sound_p / row.file_name)
    except:
        return

    n_fft = 1024
    win_length = None
    hop_length = 512
    n_mels = 128

    mel = T.MelSpectrogram(
        sample_rate  = sr,
        n_fft        = n_fft,
        win_length   = win_length,
        hop_length   = hop_length,
        center       = True,
        pad_mode     = "reflect",
        power        = 2.0,
        norm         = "slaney",
        onesided     = True,
        n_mels       = n_mels,
        mel_scale    = "htk"
    )

    melspec = mel(wav)[0]

    height = 128 * 2
    width = height * 4
    dpi = 100
    
    fig = plt.figure(frameon=False, figsize=(width/dpi, height/dpi), dpi=dpi)
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)

    im = ax.imshow(librosa.power_to_db(melspec), origin="lower", aspect="auto")

    filename = spec_p / f"{row['gbifID']}.jpg"

    plt.savefig(filename)
    plt.close()

    return row

# samples = df.sample(10, random_state=10000).apply(generateSpectogram, axis=1)


# +
# Creating dask dataframe
ddf = dd.from_pandas(df, npartitions=16)

res = ddf.map_partitions(lambda df: df.apply((lambda row: generateSpectogram(row)), axis=1)).compute(scheduler=get)

# +
print(f"Had: {df.shape[0] - len(os.listdir(spec_p))} files that couldnt be opened")
current_spectograms = os.listdir(spec_p)


new_df = df[df.apply(lambda x: True if f"{x.gbifID}.jpg" in current_spectograms else False, axis=1)]

print(df.shape, new_df.shape)
# -

# Saving

new_df.to_csv(spec_p / "../full_dataset_df.csv")

StarcoderdataJupyterScriptsDedupFiltered