import os
import json
import torch
import torchaudio

if __name__ == "__main__":
    for mode in ["worst", "best"]:
        if not os.path.isdir(mode):
            os.makedirs(mode)

        with open(f"{mode}_quality.json", "r") as fp:
            results = json.load(fp)

        for idx, (filepath, result) in enumerate(results.items()):
            filename = os.path.basename(filepath).replace(".wav", "")
            x, sr = torchaudio.load(filepath)
            x /= x.abs().max()
            start_frame = x.shape[-1] // 2
            end_frame = start_frame + 262144
            x_crop = x[:, start_frame:end_frame]

            score = result["score"]

            filepath = os.path.join(
                f"{mode}", f"{idx:02d}-{normalized_score:02.2f}-{filename}.wav"
            )
            torchaudio.save(filepath, x_crop, sr)

            if idx > 25:
                break