"""Generating testing samples for model comparisons."""

import sys

sys.path.insert(0, "/home/tony/Work/neon/sunoDiff/")

import os
from suno_utils.audio import Audio
from tqdm import tqdm
import numpy as np
from suno_utils.worker.settings import s3_client
import torch
from generation import preload_models, generate, _retrieve_models
from suno_utils.utils.s3 import read_from_s3
import json

lyrics_walking_down_street = """
[intro]

[verse]
Walking down the street, feeling so alive
Got my head in the clouds, got a gleam in my eye
Every step I take, it's like a brand new start
No matter where I'm going, I'll always find my part
(oh-oh-oh)

[chorus]
Life is like a high-wire act, we're dancing in the sky
No need to worry, no need to ask why
With a little bit of courage, we can chase our dreams
No matter what comes our way, we'll always be a team
(we're unstoppable, yeah)

[outro]
"""
# dict of titie; lyrics; tags
test_exmaples = {
    "film_epic_walking": {
        "title": "film_epic_walking",
        "text": lyrics_walking_down_street,
        "tags": "orchestral film epic",
        "gen_id": "ae137e1f-7562-48f6-a11c-b125e645282c",
    },
    "case_orange": {
        "title": "case_orange",
        "text": None,
        "tags": "80s, new wave, female power",
        "gen_id": "3b1830d7-5553-4785-be0b-ffeadf7756d8",
    },
}

if __name__ == "__main__":
    print(f"working with GPU:{os.environ['CUDA_VISIBLE_DEVICES']} ")
    # diff_ckpt_path = "/app/suno/data/dpo/models/diff_vae_25_peaq_v1_nov7.pt"
    # diff_output_dir = "diff_v1"
    diff_ckpt_path = (
        "/app/suno/checkpoints/2024-11-30_09-50-23_s5271/step_3000_infer.pt"
    )
    diff_output_dir = "diff_v1_t7"
    # diff_ckpt_path = (
    #     "/app/suno/checkpoints/2024-12-03_01-17-54_s8675/step_5000_infer.pt"
    # )
    # diff_output_dir = "diff_v1_t13"

    diff_output_path = os.path.join("/home/tony/Work/gpt_samples/", diff_output_dir)
    os.makedirs(diff_output_path, exist_ok=True)
    print(f"Sending outputs to {diff_output_path}")
    CKPT_DIR = "/home/christian/code/neon/stable-audio-tools/harmonai_train/"
    _ = preload_models(
        tokenizer_filepath="/app/suno/data/dpo/models/tokenizer_60k.json",
        semantic_model_filepath="/app/suno/data/dpo/models/mert_25.pt",
        semantic_clusters_filepath="/app/suno/data/dpo/models/mert_25_2x4k.npy",
        weights_precision=torch.bfloat16,
        compile=True,
        codec_filepath="/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth",
        dit_model_filepath=diff_ckpt_path,  # step_5000_infer.pt",
    )
    models = _retrieve_models()
    model_duration_s = 30
    if models["dit_model"].ctx_len is not None:
        model_duration_s = 6 * 60
    else:
        model_duration_s = models["dit_model"].block_size // models["dit_model"].io_hz
    duration_s = 2 * 60 if model_duration_s >= 2 * 60 else 30

    def generate_audio_with_engine(info: dict):
        """Generate audio samples for a given example."""
        print(f"Generating audio for {info['title']}")
        gen_id = info["gen_id"]
        title = info["title"]
        lyrics = info["text"]
        tags = info["tags"]
        s3_filepath = f"s3://suno-data-uploads/studio/uploads/{gen_id}.npz"
        data = read_from_s3(s3_filepath, read_f=np.load)

        if "v3.0_raw" in data:
            codes = data["v3.0_raw"]
        elif "v3.5_raw" in data:
            codes = data["v3.5_raw"]
        elif "v4.0_raw" in data:
            codes = data["v4.0_raw"]
        else:
            raise ValueError("No codes found")

        text_data = read_from_s3(
            f"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json"
        )
        # print(aligned_lyrics)
        if lyrics is None:
            aligned_lyrics = json.loads(text_data)
            lyrics = "".join(w["word"] for w in aligned_lyrics if "word" in w)

        seeds = [0, 1, 2]
        sem_codes = torch.tensor(codes[:, 0]).to(torch.long)
        steps = 12
        text_cfg_coef = 2.0
        ctx_cfg_coef = 1.0

        for n in seeds:
            pred_audio = generate(
                sem_codes,
                lyrics=lyrics,
                tags=tags,
                text_cfg_coef=text_cfg_coef,
                ctx_cfg_coef=ctx_cfg_coef,
                steps=steps,
                seed=n,
            ).normalize_volume()
            mp3_output_path = os.path.join(diff_output_path, f"{title}_{n}.mp3")
            pred_audio.to_hq_mp3(mp3_output_path)
            s3_client.upload_file(
                mp3_output_path,
                "suno-data-uploads",
                f"studio/uploads/neon230817/{diff_output_dir}/{title}_{n}.mp3",
                ExtraArgs={
                    "ContentType": "audio/mp3",
                },
            )

    print("Start generating samples")
    for title, example in tqdm(test_exmaples.items()):
        generate_audio_with_engine(example)
    print("Finished generating samples. Have a good day!")