# this script will read metadata (tags) from file and connect this to audio files
# for each audio file we will extract some audio features that describe the production
# we will then write these results to a new metadata file using the id
# during training we will load tags and the audio features which are estimated by our model

import os
import json
import torchaudio
import numpy as np
import matplotlib.pyplot as plt

from suno_amp.utils import measure_octave_band_spectrum

if __name__ == "__main__":
    meta_jsonl_filepath = "/app/suno/christian/data/metas.jsonl"
    audio_dir = "/app/suno/christian/data/ytm_audio"

    os.makedirs("plots", exist_ok=True)

    count = 0
    with open(meta_jsonl_filepath, "r") as fp:
        while True:
            json_line = json.loads(fp.readline())
            for key, val in json_line.items():
                print(key, val)
                s3_filepath = json_line["s3_filepath"]
                tags = json_line["tags"]
            if not json_line:
                break

            # get local filepath
            ytid = os.path.basename(s3_filepath).split(".")[0]
            local_filepath = os.path.join(audio_dir, os.path.basename(s3_filepath))
            x, sr = torchaudio.load(local_filepath)

            x = x[:, :524288]

            bands, energies = measure_octave_band_spectrum(x.numpy(), sr)
            print(bands, energies)

            band_labels = []
            for band in bands:
                if band < 1000:
                    band_labels.append(f"{band:0.0f}")
                else:
                    band_labels.append(f"{band/1000:0.0f}k")

            fig, axs = plt.subplots()
            axs.plot(np.arange(0, len(bands)), energies)
            axs.set_xticks(np.arange(0, len(bands)))
            axs.set_xticklabels(band_labels)
            plt.grid(c="lightgray")
            plt.title(f"{tags}", fontsize=8)
            plt.savefig(f"plots/{ytid}.png", dpi=300)
            plt.close("all")

            count += 1
            if count > 1:
                break