# Description: This script is used to encode the npz files in the input_dir and save the encoded npz files in the output_dir.
# Note this is reencoding 7b and decode 7b
from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as semantic_encode,
)
from suno_utils.tasks.dac_2c_12cb import (
    preload_models as preload_codec_models_12,
    encode as codec_encode_12,
    decode as codec_decode_12,
)
import os
import numpy as np
from tqdm import tqdm
import json
import random
import time

random_seed = int((time.time() * 1000) % 100000000)
print("Random seed set to:", random_seed)
random.seed(random_seed)

semantic_ckpt_path = "/home/victor/data/models/chirp_v2/mert_25.pt"
semantic_centroids_path = "/home/victor/data/models/chirp_v2/mert_25_2x4k.npy"
codec_ckpt_path_12 = "/app/suno/tony/v3/dac_2c_25x12.pt"

input_dir = "/app/suno/data/dpo/v3_npz"
output_dir = "/app/suno/data/dpo/7b_recycle_npz"

if __name__ == "__main__":
    avialbe_device = f"cuda:{os.environ['CUDA_VISIBLE_DEVICES']}"
    print(avialbe_device)
    preload_codec_models_12(codec_ckpt_path_12, device="cuda")
    preload_semantic_models(semantic_ckpt_path, semantic_centroids_path, device="cuda")
    print("finished loading models")

    # with open("/home/tony/Data/Preference/7b_v2/pre_model_20240412_recut_id.json", "r") as fp:
    with open("/home/tony/Data/Preference/7b_v2/interesting_clips_20240421_prev_model.json", "r") as fp:
        total_jobs = json.load(fp)
    total_jobs = [f"{file_name}.npz" for file_name in total_jobs]
    finished_jobs = os.listdir(output_dir)
    unfinished_jobs = sorted(list(set(total_jobs) - set(finished_jobs)))
    # for testing
    # total_jobs = [
    #     "fff507cb-b970-4583-8cc3-d937dbf0e4b9.npz",
    #     "fff82f8b-af8c-484e-a4e7-3bf38ef4c410.npz",
    # ]
    print(f"{len(unfinished_jobs)} to be converted.")
    # random this
    random.shuffle(unfinished_jobs)
    for filename in tqdm(unfinished_jobs):
        try:
            input_path = os.path.join(input_dir, filename)
            output_path = os.path.join(output_dir, filename)
            if os.path.exists(output_path):
                continue
            with open(output_path, "w") as fp:
                fp.write("")
            # print(input_path)
            # encode
            npz = np.load(input_path).get("v3.0_raw")
            if npz is None:
                print(f"issue with {output_path}")
                os.remove(output_path)
                continue
            codec_labels = npz[:, 1:]
            audio = codec_decode_12(codec_labels)
            # only encode 1 since we use only 1 now
            semantic_labels = semantic_encode(audio, n_codebooks=1)
            codec_labels = codec_encode_12(audio)
            n_frames = min(semantic_labels.shape[0], codec_labels.shape[0])
            audio_arr = np.concatenate(
                [semantic_labels[:n_frames, :], codec_labels[:n_frames, :]],
                axis=-1,
            )
            # print(audio_arr.shape)
            output_npz = {"v3.0_raw": audio_arr}
            np.savez(output_path, **output_npz)
        except Exception as e:
            os.remove(output_path)
            print(f"WTF {filename}, {e}")
    print("DONE!")