import glob, os, random
from concurrent.futures import ProcessPoolExecutor
import numpy as np
import soundfile as sf
SR = 16000
NOISES = sorted(glob.glob("/scratch/noise/*.flac"))
os.makedirs("/scratch/rehearsal_deg/audio", exist_ok=True)

def degrade(src):
    rng = random.Random(hash(os.path.basename(src)) & 0xffffffff)
    dst = "/scratch/rehearsal_deg/audio/" + os.path.basename(src)
    if os.path.exists(dst):
        return
    wav, _ = sf.read(src, dtype="float32")
    if rng.random() < 0.7 and NOISES:
        n, _ = sf.read(NOISES[rng.randrange(len(NOISES))], dtype="float32")
        if len(n) < len(wav):
            n = np.tile(n, len(wav) // len(n) + 1)
        s = rng.randrange(max(len(n) - len(wav), 1))
        n = n[s:s + len(wav)]
        snr = rng.uniform(3.0, 15.0)
        pw, pn = float((wav ** 2).mean()) + 1e-9, float((n ** 2).mean()) + 1e-9
        wav = wav + n * np.sqrt(pw / (pn * 10 ** (snr / 10)))
    if rng.random() < 0.5:
        f = rng.uniform(0.9, 1.1)
        idx = np.arange(0, len(wav) - 1, f)
        wav = np.interp(idx, np.arange(len(wav)), wav).astype(np.float32)
    if rng.random() < 0.5:
        idx8 = np.arange(0, len(wav) - 1, 2.0)
        down = np.interp(idx8, np.arange(len(wav)), wav)
        idx16 = np.arange(0, len(down) - 1, 0.5)
        wav = np.interp(idx16, np.arange(len(down)), down).astype(np.float32)
    wav = np.clip(wav * (10 ** (rng.uniform(-6, 3) / 20)), -1, 1).astype(np.float32)
    sf.write(dst, wav, SR)

files = sorted(glob.glob("/scratch/rehearsal/audio/*.flac"))
with ProcessPoolExecutor(max_workers=8) as ex:
    list(ex.map(degrade, files, chunksize=64))
print("DEGRADE_DONE", len(files))