import os
import argparse

from suno_utils.audio import Audio
from suno_utils.tasks import dac
from multiprocessing import Pool
from tqdm.contrib.concurrent import process_map, thread_map
from time import time
import numpy as np

SAMPLE_RATE = 24_000
EMBEDDING_RATE = 25
N_CODEBOOKS = 8
BATCH_SIZE = 100000

OUT_DATA_DIR = "/app/suno/data/mert_25hz_short"
OUT_AUDIO_DIR = os.path.join(OUT_DATA_DIR, "audio")
OUT_AUDIO_DEMUC_DIR = os.path.join(OUT_DATA_DIR, "audio_demuc")
OUT_AUDIO_DEMUC_LABEL_DIR = os.path.join(OUT_DATA_DIR, "audio_demuc_label")
OUT_TSV_DIR = os.path.join(OUT_DATA_DIR, "audio_tsv")
OUT_LABEL_DIR = os.path.join(OUT_DATA_DIR, "label")
OUT_TEMP_DIR = os.path.join(OUT_DATA_DIR, "temp")
# CODEC_CKPT_PATH = "s3://suno-data/georg/trained_models/chirp_v1/codec.pt"
CODEC_CKPT_PATH = "/home/tony/Data/MERT/codec.pt"
MAX_SHAPE = 300


def merge_np_data(file_names):
    i = 0
    max_shape = 300
    data_split = "train"
    output_array = np.zeros((len(file_names), max_shape), dtype=np.int16)
    for i, test_file_name in enumerate(file_names):
        test_id = test_file_name.replace(".wav", "")
        file_to_fetch = f"/app/suno/data/mert_25hz_short/audio_mert_label/{data_split}_{test_id}.npy"
        if not os.path.exists(file_to_fetch):
            print("WTF", file_to_fetch)
            break
        try:
            test_outputs = np.load(file_to_fetch, allow_pickle=True)
            test_outputs = test_outputs.reshape(-1)
            if test_outputs.shape[0] < max_shape:
                test_outputs = np.pad(
                    test_outputs,
                    (0, max_shape - test_outputs.shape[0]),
                    "constant",
                    constant_values=-1,
                )
            output_array[i, :] = test_outputs
        except Exception as E:
            print(i, test_id, E)
            output_array[i, :] = np.ones(300) * -1
    print(i)
    output_array = output_array.astype(np.int16)
    np.save(
        os.path.join(OUT_LABEL_DIR, f"{data_split}.mert_{file_names[0]}.npy"),
        output_array,
    )
    return


if __name__ == "__main__":
    start_time = time()
    # load current data
    tsv_info = []
    # check GPU
    with open(os.path.join(OUT_TSV_DIR, "train.tsv"), "r") as f:
        for line in f.read().strip().split("\n"):
            if len(line.strip()) == 0:
                continue
            tsv_info.append(line.strip().split("\t"))
    tsv_data_path = tsv_info[0]
    tsv_data = tsv_info[1:]
    print(len(tsv_data))
    work_items = [x[0] for x in tsv_data]
    batches = [
        work_items[x : x + BATCH_SIZE] for x in range(0, len(work_items), BATCH_SIZE)
    ]
    print(len(work_items), "work_items", len(batches), "batches")
    thread_map(merge_np_data, batches, max_workers=20, chunksize=1)
    # for batch in batches:
    #     encode_audio_file(batch)
    print(f"Done~!, {time() - start_time:.3f} seconds.")