from suno_utils.tasks.dac_2c import (
    preload_models as preload_codec_models_8,
    decode as codec_decode_8,
)
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,
)
import os
import numpy as np
from tqdm import tqdm


codec_ckpt_path_8 = "/home/victor/data/models/chirp_v2/dac_2c_25x8.pt"
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/npz"
output_dir = "/app/suno/data/dpo/7b_npz"

if __name__ == "__main__":
    avialbe_device = f"cuda:{os.environ['CUDA_VISIBLE_DEVICES']}"
    print(avialbe_device)
    preload_codec_models_8(codec_ckpt_path_8, device="cuda")
    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")

    while True:
        total_jobs = os.listdir(input_dir)
        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",
        # ]
        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("v1_raw")
                if npz is None:
                    # this is arleady a 7b...
                    np.save(npz, output_path)
                    continue
                codec_labels = npz[:, 1:]
                audio = codec_decode_8(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 = {"v2_raw": audio_arr}
                np.savez(output_path, **output_npz)
            except Exception as e:
                print(f"WTF {filename}, {e}")
        new_total_jobs = os.listdir(input_dir)
        if total_jobs != new_total_jobs:
            print("More downloaded jobs found! Repeat the process")
        else:
            print("All downloaded jobs are processed!")
            break
    print("DONE!")