# 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,
)
from suno_utils.fine.generation import FineConfig, Fine, load_model as load_fine_modal, generate
from suno_utils.models.dac.model.dac2_12 import DAC as DAC12
import torch
import os
import numpy as np
from tqdm import tqdm
import json
import random
import typing as tp
from suno_utils.audio import Audio
import torch.nn.functional as F

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"
fine_codec_ckpt_path = "/home/victor/data/models/dac_100x16.pth"
fine_ckpt_path = "/app/suno/checkpoints/7b_fine/best_ckpt.pt"  # 7b

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

class FineCodec(torch.nn.Module):
    def __init__(self, codec_path: str):
        super().__init__()

        sd = torch.load(codec_path, map_location="cpu")
        model = DAC12(**sd["metadata"]["kwargs"])
        model.load_state_dict(sd["state_dict"])
        model.eval()
        self.model = model

    @torch.inference_mode()
    def forward(
        self,
        audios: tp.Union[torch.Tensor, tp.List[torch.Tensor], tp.Tuple[torch.Tensor]],
        device: tp.Union[torch.device, str],
    ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
        if isinstance(audios, list) or isinstance(audios, tuple):
            audios = torch.cat(audios, dim=0)  # (B, C, T)

        audios = audios.to(device)
        if len(audios.shape) == 2:
            audios = audios.unsqueeze(0)
        encoded = self.model(audios, 48000)
        return encoded["codes"], encoded["z"], encoded["audio"]

    @torch.inference_mode()
    def decode(self, codes):
        # zero out padding for now
        codes = torch.where(codes < self.model.quantizer.codebook_size, codes, 0)
        latents = self.model.quantizer.from_codes(codes)[0]
        return self.model.decode(latents)


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")
    fine_codec = FineCodec(fine_codec_ckpt_path).cuda()
    fine_model = load_fine_modal(fine_ckpt_path, "cuda")
    config = fine_model.config
    print("finished loading models")

    # setup some helper functions inside
    def pad_coarse_codes(x, remove_last=0):
        if remove_last > 0:
            # replace with padding
            x[:, -remove_last:] = config.coarse_pad_token
        # pad to length + 1 for inference token
        x = F.pad(
            x,
            (0, config.coarse_samples - x.shape[-1] + 1),
            "constant",
            config.coarse_pad_token,
        )

        # add fine streams
        x = F.pad(
            x,
            (0, 0, 0, config.fine_n_codebooks),
            "constant",
            config.fine_pad_token,
        )
        # set inference token
        x[:, config.coarse_n_codebooks :, -1] = config.fine_infer_token
        return x.squeeze(0)

    def upsample_chirp(filename: str):
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename)
        if os.path.exists(output_path):
            return
        with open(output_path, "w") as fp:
            fp.write("")
        try:
            # encode
            npz = np.load(input_path).get("v3.0_raw")
            array = torch.from_numpy(npz).long().to("cuda").T
            original_duration = array.shape[1] / 25
            # make coarse batch by splitting into 2s clips
            coarse_codes = []
            for i in range(0, array.size(-1), config.coarse_samples):
                code = array[1:, i : i + config.coarse_samples]
                coarse_codes.append(pad_coarse_codes(code.unsqueeze(0)))
            coarse_codes = torch.stack(coarse_codes, dim=0)
            # print(coarse_codes.shape)
            # print(coarse_codes.max())
            # this should be somewhat deterministic
            fine_codes = generate(
                fine_model, coarse_codes, config.t_fine - 1, temperature=0.8, top_k=5, silent=True
            )
            fine_codes = torch.transpose(fine_codes, 0, 1)
            fine_codes = fine_codes.reshape(1, config.fine_n_codebooks, -1)

            audio = fine_codec.decode(fine_codes.cuda())
            audio = Audio.from_array_float(audio.squeeze(0).cpu().numpy(), 48000)
            audio = audio.get_segment(to_s=original_duration)
            
            # 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,
            )
            output_npz = {"v3.0_raw": audio_arr}
            np.savez(output_path, **output_npz)
            return 
        except Exception as e:
            os.remove(output_path)
            print(f"WTF {filename}, {e}")


    with open("/home/tony/Data/Preference/7b_v2/7v_v20_full_recut_id_negative.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):
        upsample_chirp(filename)
    print("DONE!")