"""Generating testing samples for model comparisons.

aws s3 cp audio_files.json s3://suno-data-uploads/studio/uploads/neon230817/audio_files.json
aws cloudfront create-invalidation --distribution-id E3ROE628KJYDYL --paths "/neon230817/audio_files.json"
"""

import os
import torch
from suno_utils.audio import Audio
from suno_utils.gpt.chirp_v2_5 import (
    GenerationConfig,
    decode_stream,
    preload_models,
    prep_gconf,
)
from suno_utils.gpt.generation_engine import (
    align_codes,
    make_request,
    unshift_arrays_v2,
)
from suno_utils.gpt.engine import Engine
from tqdm import tqdm
import numpy as np
from suno_utils.worker.settings import s3_client


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]
"""
lyrics_chinese_pop_woman = """
[Verse]
一盏离愁孤灯伫立在窗口
我在门后假装你人还没走
旧地如重游月圆更寂寞
夜半清醒的烛火不忍苛责我

[Verse]
一壶漂泊浪迹天涯难入喉
你走之后酒暖回忆思念瘦
水向东流时间怎么偷
花开就一次成熟我却错过

[Chorus]
谁在用琵琶弹奏一曲东风破
岁月在墙上剥落看见小时候
犹记得那年我们都还很年幼
而如今琴声幽幽我的等候你没听过

[outro]
"""

lyrics_ft = """
[Intro]
[Verse]
I used to think it's all a game
That somethin' more was to blame
But now I see, it's in the details, baby
Just a little tweak here and there
Can make our model beyond compare
Yeah, it's the small things that drive it crazy (ooh-yeah)

[Pre-Chorus]
Needed to get rid of the rest,
Cause baby you deserve what's best
Gotta give you the cream of the crop
So you can get us back to the top

[Chorus]
Fine tuning, baby, it's the key
Turnin' up the signal for you and me
When we're off, we find our way back again (oh-oh)
Fine tuning, baby, it's the truth
We make it work, we're bulletproof
Together we train a model that's the best (oh-oh-oh-oh)

[end]
"""


lyrics_belong_with_me = """
[Intro]


[Verse 1]
You're on the phone with your girlfriend, she's upset
She's going off about something that you said
'Cause she doesn't get your humor like I do
I'm in the room, it's a typical Tuesday night
I'm listening to the kind of music she doesn't like
And she'll never know your story like I do

[Pre-Chorus]
But she wears short skirts, I wear T-shirts
She's Cheer Captain and I'm on the bleachers
Dreaming 'bout the day when you wake up and find
That what you're looking for has been here the whole time

[Chorus]
If you could see that I'm the one who understands you
Been here all along, so why can't you see?
You belong with me, you belong with me

[Verse 2]
Walking the streets with you and your worn-out jeans
I can't help thinking this is how it ought to be
Laughing on a park bench, thinking to myself
Hey, isn't this easy?
And you've got a smile that could light up this whole town
I haven't seen it in a while since she brought you down
You say you're fine, I know you better than that
Hey, whatcha doing with a girl like that?

[Pre-Chorus]
She wears high heels, I wear sneakers
She's Cheer Captain and I'm on the bleachers
Dreaming 'bout the day when you wake up and find
That what you're looking for has been here the whole time

[Chorus]
If you could see that I'm the one who understands you
Been here all along, so why can't you see?
You belong with me
Standing by and waiting at your back door
All this time, how could you not know, baby?
You belong with me, you belong with me

[Bridge]
Oh, I remember you driving to my house
In the middle of the night
I'm the one who makes you laugh
When you know you're 'bout to cry
And I know your favorite songs
And you tell me 'bout your dreams
Think I know where you belong
Think I know it's with me

[Chorus]
Can't you see that I'm the one who understands you
Been here all along, so why can't you see?
You belong with me
Standing by and waiting at your back door
All this time, how could you not know, baby?
You belong with me, you belong with me

[Outro]
You belong with me
Have you ever thought just maybe
You belong with me?
You belong with me
"""
lyrics_yum_yum = """
[verse]
A is for the amazing grace that we recieve
B is for the blessings, every day we blieve
C is for the chorus, we sing it loud and clear
D is for the devotion that we hold dear

[Chorus]
B is for Buttocks, ripe and slightly damp
Yum Yum, boy oh boy do I like God
Every day is a gift, when you won a skateboard ramp.
Yum yum, boy oh boy do I like God.

[outro]
"""

lyrics_country_road = """
[Verse 1]
Almost Heaven, West Virginia
Blue Ridge Mountains, Shenandoah River
Life is old there, older than the trees
Younger than the mountains, growing like a breeze

[Chorus]
Country roads, take me home
To the place I belong
West Virginia, mountain mama
Take me home, country roads
"""

lyrics_rap = """
[Verse]
I've created a monster
'Cause nobody wants to see Marshall no more, they want Shady, I'm chopped liver
Well, if you want Shady, this is what I'll give ya
A little bit of weed mixed with some hard liquor
Some vodka that'll jump-start my heart quicker
Than a shock when I get shocked at the hospital
By the doctor when I'm not cooperating
When I'm rockin' the table while he's operating (Hey!)
You waited this long, now stop debating
'Cause I'm back, I'm on the rag and ovulating
I know that you got a job, Ms. Cheney
But your husband's heart problem's complicating
So the FCC won't let me be
Or let me be me, so let me see
They tried to shut me down on MTV
But it feels so empty without me
So come on and dip, bum on your lips
Fuck that, cum on your lips and some on your tits
And get ready, 'cause this shit's about to get heavy
I just settled all my lawsuits (Fuck you, Debbie!)

[Chorus]
Now, this looks like a job for me
So everybody, just follow me
'Cause we need a little controversy
'Cause it feels so empty without me
I said this looks like a job for me
So everybody, just follow me
'Cause we need a little controversy
'Cause it feels so empty without me

[end]
"""

# dict of titie; lyrics; tags
test_exmaples = {
    # "chinese_pop_woman": {
    #     "text": lyrics_chinese_pop_woman,
    #     "tags": "pop woman",
    # },
    # "film_epic_walking": {
    #     "text": lyrics_walking_down_street,
    #     "tags": "orchestral film epic",
    # },
    # "ft_kpop_female": {
    #     "text": lyrics_ft,
    #     "tags": "k-pop female",
    # },
    # "country_road_bluegrass_female": {
    #     "text": lyrics_country_road,
    #     "tags": "bluegrass, female",
    # },
    # "yum_chant": {
    #     "text": lyrics_yum_yum,
    #     "tags": "gregorian chant",
    # },
    "ts_pop": {
        "text": lyrics_belong_with_me,
        "tags": "pop country, female vocal, fast tempo, taylor swift",
    },
    # "film_epic": {
    #     "text": "",
    #     "tags": "orchestral film epic",
    # },
    # "yum_gospel": {
    #     "text": lyrics_yum_yum,
    #     "tags": "gospel soulful choir",
    # },
    # "fast_female_walking": {
    #     "text": lyrics_walking_down_street,
    #     "tags": "fast female vocalist, pop-punk, pop, melodic, hd stems, goes hard, 155bpm",
    # },
    # "rap_female": {
    #     "text": lyrics_rap,
    #     "tags": "rap, female",
    # },
    # "synthwave": {
    #     "text": "",
    #     "tags": "playful energetic synthwave",
    # },
    # "beethoven_pc": {
    #     "text": "",
    #     "tags": "piano concerto, beethoven",
    # },
    # "beethoven_sym": {
    #     "text": "",
    #     "tags": "orchestral symphony, beethoven",
    # },
    # "bluegrass_instrumental": {
    #     "text": "",
    #     "tags": "bluegrass instrumental",
    # },
}

if __name__ == "__main__":
    print(f"working with GPU:{os.environ['CUDA_VISIBLE_DEVICES']} ")
    # gpt_ckpt_path = "/app/suno/checkpoints/2024-08-23_13-28-02/last_ckpt_infer.pt"
    # gpt_output_path = os.path.join("/home/tony/Work/gpt_samples/", "7b_delay_test")
    # gpt_ckpt_path = "/app/suno/checkpoints/2024-06-17_01-25-46/last_ckpt_infer.pt"
    # gpt_output_dir = "13b_s8_redo"
    # gpt_ckpt_path = "/app/suno/checkpoints/2024-08-17_22-00-35/last_ckpt_infer.pt"
    # gpt_output_path = os.path.join(
    #     "/home/tony/Work/gpt_samples/", "13b_2h_multi_ft_t2_10k_short_fast"
    # )
    # gpt_ckpt_path = "/app/suno/data/dpo/models/model_30b_fix_ft2_20k.pt"
    # gpt_output_dir = "30b_t1_jul1"
    # gpt_ckpt_path = "/app/suno/checkpoints/2024-08-07_18-25-41/last_ckpt_infer.pt"
    # gpt_output_path = os.path.join("/home/tony/Work/gpt_samples/", "30b_t1_v16_sft_raw")
    gpt_ckpt_path = "/app/suno/checkpoints/2024-08-29_04-23-05/last_ckpt_infer.pt"
    gpt_output_dir = "30b_t3_v1_repro7"
    gpt_output_path = os.path.join("/home/tony/Work/gpt_samples/", gpt_output_dir)
    # gpt_ckpt_path = "/app/suno/checkpoints/2024-08-30_01-33-26/last_ckpt_infer.pt"
    # gpt_output_path = os.path.join("/home/tony/Work/gpt_samples/", "30b_t3_v2")
    os.makedirs(gpt_output_path, exist_ok=True)
    print(f"Sending outputs to {gpt_output_path}")
    MAX_STREAMS = 4 if "30b" not in gpt_output_path else 3
    MAX_SEQ = (
        32 if "30b" not in gpt_output_path else 12
    )  # 13b 8 head max is 40 and OOM...

    # Import models and setup
    preload_models(load_gpt=False, load_semantic=False, load_codec_device="cuda")

    engine = Engine(
        gpt_ckpt_path,
        "/app/suno/data/dpo/models/tokenizer_60k.json",
        max_sequences=MAX_SEQ,
        compile=False,
    )

    def generate_audio_with_engine(text, text_tags, title, random_seed=42):
        general_config = dict(
            cfg_coef=1.0,  # no text cfg for dpo stream
            min_eos_p=0.1,
            eos_pad_duration_s=0,
            # cfg_coef_tags=0.0,
            cfg_coef_tags_max_steps=None,  # collect the data for now
            cfg_coef_tags=2,
            cfg_coef_neg_tags=-1,
            text_neg_tags="repetitive, loop"
            if "30b" in gpt_output_path
            else "repetitive, loop, noisy, distorted",
            n_repeat_tags=1,
            use_whisper=False,
            text_start_control_tags="{start:0} "
            if "30b" in gpt_output_path
            else "{start} ",
            # text_end_control_tags="{end}",
            min_text_offset=0,  # this is required for 30b
            random_seed=random_seed,
            n_batch=1,
        )

        if "short" in gpt_output_path:
            general_config["text_start_control_tags"] = "{start:0;vocals:start}"
            general_config["text_end_control_tags"] = "{end}"
        cfg = GenerationConfig(
            text=text, text_tags=text_tags, max_gen_duration_s=125, **general_config
        )
        preped_cfg = prep_gconf(cfg)
        model_conf = engine.model.config
        requests = []
        for i in range(1, MAX_STREAMS + 1):
            request = make_request(
                f"{i}", preped_cfg, engine.model.config, engine.tokenizer
            )
            requests.append(request)

        jobs = engine.run_request(requests, tqdm_enabled=True)

        for i, job in enumerate(jobs):
            stream = engine.token_generator(job)
            audios = []
            codes = []
            raw_codes = []
            tensor_codes = []
            for code in stream:
                raw_codes.append(code)
            torch_codes = torch.stack(raw_codes).cpu()
            # print(torch_codes.shape)
            unshifted_codes = unshift_arrays_v2(torch_codes, model_conf)
            # print(unshifted_codes.shape)
            for code in align_codes(tqdm(raw_codes), model_conf):
                codes.append(code.numpy().astype(np.int16))
                tensor_codes.append(code)
            for audio in decode_stream(tensor_codes):
                audios.append(audio[0])
            with open(os.path.join(gpt_output_path, f"{title}_{i}.npz"), "wb") as f:
                np.savez(f, codes=unshifted_codes)
            audio_continued = Audio.concatenate(audios)
            # audio_continued.play()
            # save the file to mp3
            mp3_output_path = os.path.join(
                gpt_output_path, f"{title}_{i}_seed{random_seed}.mp3"
            )
            audio_continued.to_hq_mp3(mp3_output_path)
            s3_client.upload_file(
                mp3_output_path,
                "suno-data-uploads",
                f"studio/uploads/neon230817/{gpt_output_dir}/{title}_{i}_seed{random_seed}.mp3",
                ExtraArgs={
                    "ContentType": "audio/mp3",
                },
            )
        return

    print("Start generating samples")
    for random_seed in range(42, 43):
        for title, example in tqdm(test_exmaples.items()):
            generate_audio_with_engine(
                example["text"], example["tags"], title=title, random_seed=random_seed
            )
    print("Finished generating samples. Have a good day!")
