from audiomentations import ApplyImpulseResponse, AddGaussianNoise, TimeStretch, PitchShift, Compose
import itertools
import numpy as np
augmentations = {
    "timestretch": TimeStretch(min_rate=0.5, max_rate=2.0, p=1.0, leave_length_unchanged=False),
    "pitchshift": PitchShift(min_semitones=-8, max_semitones=8, p=1.0),
    "gaussian": AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.015, p=1.0),
    "rir": ApplyImpulseResponse(ir_path="/home/joregan/.psst/rir", p=1.0),
}
/home/joregan/miniconda3/envs/fs2/lib/python3.8/site-packages/audiomentations/augmentations/apply_impulse_response.py:48: UserWarning: The default value of leave_length_unchanged will change from False to True in a future version of audiomentations. You can set the value explicitly to remove this warning for now.
  warnings.warn(
all2=list(itertools.combinations_with_replacement(augmentations.keys(), 2))
all3=list(itertools.combinations_with_replacement(augmentations.keys(), 3))
all2 = [a for a in all2 if a[0] != a[1]]
all3 = [a for a in all3 if(a[0] != a[1] and a[1] != a[2] and a[0] != a[2])]
combinations = {}
for item in all2:
    key = "-".join(item)
    val = Compose([augmentations[item[0]], augmentations[item[1]]])
    combinations[key] = val
for item in all3:
    key = "-".join(item)
    val = Compose([augmentations[item[0]], augmentations[item[1]], augmentations[item[2]]])
    combinations[key] = val
combinations["-".join(augmentations.keys())] = Compose([v for v in augmentations.values()])
combinations
{'timestretch-pitchshift': <audiomentations.core.composition.Compose at 0x7fc658e47100>,
 'timestretch-gaussian': <audiomentations.core.composition.Compose at 0x7fc658e47f10>,
 'timestretch-rir': <audiomentations.core.composition.Compose at 0x7fc658e47fa0>,
 'pitchshift-gaussian': <audiomentations.core.composition.Compose at 0x7fc658e27430>,
 'pitchshift-rir': <audiomentations.core.composition.Compose at 0x7fc658e27880>,
 'gaussian-rir': <audiomentations.core.composition.Compose at 0x7fc658e272b0>,
 'timestretch-pitchshift-gaussian': <audiomentations.core.composition.Compose at 0x7fc658e47970>,
 'timestretch-pitchshift-rir': <audiomentations.core.composition.Compose at 0x7fc658e275b0>,
 'timestretch-gaussian-rir': <audiomentations.core.composition.Compose at 0x7fc658e27340>,
 'pitchshift-gaussian-rir': <audiomentations.core.composition.Compose at 0x7fc658e27160>,
 'timestretch-pitchshift-gaussian-rir': <audiomentations.core.composition.Compose at 0x7fc65bee4190>}
_DIR = "/home/joregan/psst-data/psst-data-2022-03-02/train"
_TSV = "/home/joregan/.psst/out/data/psst-fairseq/train.tsv"
TIMITFILES = []
with open(_TSV) as f:
    for line in f.readlines():
        if line.startswith("TEST/") or line.startswith("TRAIN/"):
            parts = line.split("\t")
            TIMITFILES.append(parts[0])
from pathlib import Path
_DIR_PATH = Path(_DIR)
import soundfile as sf
check = [a.replace("TRAIN/", "").replace("TEST/", "") for a in TIMITFILES]
assert len(check) == len(set(check))
all = {**augmentations, **combinations}
_TSV_DIR = Path("/home/joregan/.psst/out/data/psst-fairseq/")
COMMON_PATH = "/home/joregan/psst-data/psst-data-2022-03-02/train/"
for file in TIMITFILES:
    if not (file.startswith("TRAIN/") or file.startswith("TEST/")):
        continue
    filepath = _DIR_PATH / file
    filestr = str(filepath)
    raw_audio, sr = sf.read(filestr)
    orig_audio = np.array(raw_audio)
    for aug in all.keys():
        out_tsv = _TSV_DIR / f"{str(aug)}.tsv"
        if "/TRAIN/" in filestr:
            outfile = filestr.replace("/TRAIN/", f"/{aug}/")
        elif "/TEST/" in filestr:
            outfile = filestr.replace("/TEST/", f"/{aug}/")
        else:
            print(f"Error reading file: {filestr}")
            continue
        outpath = Path(outfile)
        outpath.parent.mkdir(parents=True, exist_ok=True)
        augmenter = all[aug]
        try:
            newaudio = augmenter(samples=orig_audio, sample_rate=sr)
        except AttributeError:
            raise AttributeError("Error with augmentation", aug)
        out_frames = len(newaudio)
        writable_filestr = outfile.replace(COMMON_PATH, "")
        with open(out_tsv, "a") as tsvf:
            tsvf.write(f"{writable_filestr}\t{out_frames}\n")
        sf.write(outfile, newaudio, sr)
all_the_things = """
gaussian-rir
gaussian
pitchshift-gaussian-rir
pitchshift-gaussian
pitchshift-rir
pitchshift
rir
timestretch-gaussian-rir
timestretch-gaussian
timestretch-pitchshift-gaussian-rir
timestretch-pitchshift-gaussian
timestretch-pitchshift-rir
timestretch-pitchshift
timestretch-rir
timestretch
"""
%%writefile runner.sh
for i in gaussian-rir \
gaussian \
pitchshift-gaussian-rir \
pitchshift-gaussian \
pitchshift-rir \
pitchshift \
rir \
timestretch-gaussian-rir \
timestretch-gaussian \
timestretch-pitchshift-gaussian-rir \
timestretch-pitchshift-gaussian \
timestretch-pitchshift-rir \
timestretch-pitchshift \
timestretch-rir \
timestretch; do
    #mkdir $i
    cp psst-fairseq/*.ltr $i/
    cp psst-fairseq/test.tsv psst-fairseq/valid.tsv $i/
    head -n 2299 psst-fairseq/train.tsv > $i/train.tsv
    cat psst-fairseq/$i.tsv |sed -e 's#/home/joregan/psst-data/psst-data-2022-03-02/train/##' >> $i/train.tsv
done
import torch
import glob
for model in glob.glob("out/models/**/*last.pt"):
    mymodel = torch.load(model, map_location=torch.device('cpu'))
    print(model, mymodel['args'].data)
/home/joregan/miniconda3/envs/fs2/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
out/models/psst-04_06_2022_21_18/checkpoint_last.pt ./out/data/gaussian-rir
out/models/psst-04_07_2022_08_45/checkpoint_last.pt ./out/data/timestretch-pitchshift
out/models/psst-04_05_2022_17_43/checkpoint_last.pt ./out/data/psst-fairseq
out/models/psst-04_06_2022_21_26/checkpoint_last.pt ./out/data/timestretch-pitchshift
out/models/psst-04_05_2022_16_00/checkpoint_last.pt ./out/data/psst-fairseq
out/models/psst-04_07_2022_00_14/checkpoint_last.pt ./out/data/timestretch-rir
out/models/psst-04_07_2022_08_43/checkpoint_last.pt ./out/data/gaussian-rir
out/models/psst-04_07_2022_00_10/checkpoint_last.pt ./out/data/pitchshift
out/models/psst-04_07_2022_08_44/checkpoint_last.pt ./out/data/timestretch-pitchshift-gaussian-rir
out/models/psst-04_06_2022_21_23/checkpoint_last.pt ./out/data/rir
out/models/psst-04_07_2022_00_17/checkpoint_last.pt ./out/data/timestretch-pitchshift-rir
out/models/psst-04_06_2022_21_22/checkpoint_last.pt ./out/data/pitchshift-rir
out/models/psst-baseline/checkpoint_last.pt ./out/data/timestretch-pitchshift
out/models/psst-04_06_2022_21_24/checkpoint_last.pt ./out/data/timestretch-gaussian
out/models/psst-04_05_2022_16_43/checkpoint_last.pt ./out/data/psst-fairseq
out/models/psst-04_07_2022_07_39/checkpoint_last.pt ./out/data/pitchshift-gaussian
out/models/psst-04_07_2022_00_09/checkpoint_last.pt ./out/data/timestretch-gaussian-rir
out/models/psst-04_07_2022_00_04/checkpoint_last.pt ./out/data/gaussian
out/models/psst-04_05_2022_15_43/checkpoint_last.pt ./out/data/psst-fairseq
out/models/psst-04_06_2022_21_21/checkpoint_last.pt ./out/data/pitchshift-gaussian-rir
out/models/psst-04_06_2022_21_25/checkpoint_last.pt ./out/data/timestretch-pitchshift-gaussian
out/models/psst-04_06_2022_21_15/checkpoint_last.pt ./out/data/timestretch