import gradio as gr
import subprocess
import os
import shutil
import tempfile
import spaces
import torch
import sys
import uuid
import re

print("Installing flash-attn...")
# Install flash attention
subprocess.run(
    "pip install flash-attn --no-build-isolation",
    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
    shell=True,
)

from huggingface_hub import snapshot_download

# Create xcodec_mini_infer folder
folder_path = './xcodec_mini_infer'

# Create the folder if it doesn't exist
if not os.path.exists(folder_path):
    os.mkdir(folder_path)
    print(f"Folder created at: {folder_path}")
else:
    print(f"Folder already exists at: {folder_path}")

snapshot_download(
    repo_id="m-a-p/xcodec_mini_infer",
    local_dir="./xcodec_mini_infer"
)

# Change to the "inference" directory
inference_dir = "."
try:
    os.chdir(inference_dir)
    print(f"Changed working directory to: {os.getcwd()}")
except FileNotFoundError:
    print(f"Directory not found: {inference_dir}")
    exit(1)

sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))


# don't change above code

import argparse
import numpy as np
import json
from omegaconf import OmegaConf
import torchaudio
from torchaudio.transforms import Resample
import soundfile as sf
from tqdm import tqdm
from einops import rearrange
from codecmanipulator import CodecManipulator
from mmtokenizer import _MMSentencePieceTokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
from models.soundstream_hubert_new import SoundStream
from vocoder import build_codec_model, process_audio
from post_process_audio import replace_low_freq_with_energy_matched

# Initialize device
device = "cuda:0"

# Load models once and reuse
print("Loading models...")
model = AutoModelForCausalLM.from_pretrained(
    "m-a-p/YuE-s1-7B-anneal-en-cot",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
).to(device).eval()

basic_model_config = './xcodec_mini_infer/final_ckpt/config.yaml'
resume_path = './xcodec_mini_infer/final_ckpt/ckpt_00360000.pth'
config_path = './xcodec_mini_infer/decoders/config.yaml'
vocal_decoder_path = './xcodec_mini_infer/decoders/decoder_131000.pth'
inst_decoder_path = './xcodec_mini_infer/decoders/decoder_151000.pth'

# Load codec model
model_config = OmegaConf.load(basic_model_config)
codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
codec_model.load_state_dict(torch.load(resume_path, map_location='cpu')['codec_model'])
codec_model.eval()

# Preload and compile vocoders
vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
vocal_decoder.to(device).eval()
inst_decoder.to(device).eval()

# Tokenizer and codec tool
mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
codectool = CodecManipulator("xcodec", 0, 1)

def generate_music(genre_txt, lyrics_txt, max_new_tokens=5, run_n_segments=2, use_audio_prompt=False, audio_prompt_path="", prompt_start_time=0.0, prompt_end_time=30.0, rescale=False):
    if use_audio_prompt and not audio_prompt_path:
        raise FileNotFoundError("Please provide an audio prompt filepath when enabling 'use_audio_prompt'!")

    max_new_tokens *= 100
    top_p = 0.93
    temperature = 1.0
    repetition_penalty = 1.2

    # Split lyrics into segments
    def split_lyrics(lyrics):
        pattern = r"\[(\w+)\](.*?)\n(?=\[|\Z)"
        segments = re.findall(pattern, lyrics, re.DOTALL)
        return [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]

    lyrics = split_lyrics(lyrics_txt + "\n")
    full_lyrics = "\n".join(lyrics)
    prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genre_txt.strip()}\n{full_lyrics}"] + lyrics

    raw_output = None
    stage1_output_set = []

    class BlockTokenRangeProcessor(LogitsProcessor):
        def __init__(self, start_id, end_id):
            self.blocked_token_ids = list(range(start_id, end_id))

        def __call__(self, input_ids, scores):
            scores[:, self.blocked_token_ids] = -float("inf")
            return scores
    
    for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
        section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
        guidance_scale = 1.5 if i <= 1 else 1.2

        if i == 0:
            continue

        if i == 1 and use_audio_prompt:
            audio_prompt = load_audio_mono(audio_prompt_path)
            audio_prompt = audio_prompt.unsqueeze(0).to(device)
            raw_codes = codec_model.encode(audio_prompt, target_bw=0.5).transpose(0, 1).cpu().numpy().astype(np.int16)
            audio_prompt_codec = codectool.npy2ids(raw_codes[0])[int(prompt_start_time * 50): int(prompt_end_time * 50)]
            audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
            sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
            head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
        else:
            head_id = mmtokenizer.tokenize(prompt_texts[0])

        prompt_ids = head_id + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
        prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)

        input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids

        max_context = 16384 - max_new_tokens - 1
        if input_ids.shape[-1] > max_context:
            input_ids = input_ids[:, -(max_context):]

        with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
            output_seq = model.generate(
                input_ids=input_ids,
                max_new_tokens=max_new_tokens,
                min_new_tokens=100,
                do_sample=True,
                top_p=top_p,
                temperature=temperature,
                repetition_penalty=repetition_penalty,
                eos_token_id=mmtokenizer.eoa,
                pad_token_id=mmtokenizer.eoa,
                logits_processor=LogitsProcessorList([
                    BlockTokenRangeProcessor(0, 32002),
                    BlockTokenRangeProcessor(32016, 32016)
                ]),
                guidance_scale=guidance_scale,
                use_cache=True,
                top_k=50,
                num_beams=1
            )

        if output_seq[0][-1].item() != mmtokenizer.eoa:
            tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(device)
            output_seq = torch.cat((output_seq, tensor_eoa), dim=1)

        raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1) if i > 1 else output_seq

    # Process and save outputs
    ids = raw_output[0].cpu().numpy()
    soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
    eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()

    vocals, instrumentals = [], []
    for i in range(len(soa_idx)):
        codec_ids = ids[soa_idx[i] + 1:eoa_idx[i]]
        if codec_ids[0] == 32016:
            codec_ids = codec_ids[1:]
        codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
        vocals.append(codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0]))
        instrumentals.append(codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1]))

    vocals = np.concatenate(vocals, axis=1)
    instrumentals = np.concatenate(instrumentals, axis=1)

    # Decode and mix audio
    decoded_vocals = codec_model.decode(torch.as_tensor(vocals.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device)).cpu().squeeze(0)
    decoded_instrumentals = codec_model.decode(torch.as_tensor(instrumentals.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device)).cpu().squeeze(0)

    mixed_audio = (decoded_vocals + decoded_instrumentals) / 2
    return (16000, mixed_audio.detach().numpy())

@spaces.GPU(duration=120)
def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=10):
    try:
        return generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments, max_new_tokens=max_new_tokens)
    except Exception as e:
        gr.Warning("An Error Occurred: " + str(e))
        return None

# Gradio Interface
with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation")
        gr.HTML("""
        <div style="display:flex;column-gap:4px;">
            <a href="https://github.com/multimodal-art-projection/YuE">
                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
            </a> 
            <a href="https://map-yue.github.io">
                <img src='https://img.shields.io/badge/Project-Page-green'>
            </a>
            <a href="https://huggingface.co/spaces/innova-ai/YuE-music-generator-demo?duplicate=true">
                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
            </a>
        </div>
        """)
        with gr.Row():
            with gr.Column():
                genre_txt = gr.Textbox(label="Genre")
                lyrics_txt = gr.Textbox(label="Lyrics")

            with gr.Column():
                num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
                max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=30, step=1, value=5,
                                           interactive=True)
                submit_btn = gr.Button("Submit")
                music_out = gr.Audio(label="Audio Result")

#         gr.Examples(
#             examples=[
#                 ["Rap, Hip-Hop, Street Vibes, Tough, Piercing Vocals, Piano, Synthesizer, Clear Male Vocals",
#                 """[verse]
# Woke up in the morning, sun is shining bright
# Chasing all my dreams, gotta get my mind right
# City lights are fading, but my vision's clear
# Got my team beside me, no room for fear
# Walking through the streets, beats inside my head
# Every step I take, closer to the bread
# People passing by, they don't understand
# Building up my future with my own two hands
#                 """],
#             ],
#             inputs=[genre_txt, lyrics_txt],
#             outputs=[music_out],
#             cache_examples=True,
#             cache_mode="eager",
#             fn=infer
#         )
        
        gr.Examples(
            examples=[
                [
                    "female blues airy vocal bright vocal piano sad romantic guitar jazz",
                    """[verse]
In the quiet of the evening, shadows start to fall
Whispers of the night wind echo through the hall
Lost within the silence, I hear your gentle voice
Guiding me back homeward, making my heart rejoice

[chorus]
Don't let this moment fade, hold me close tonight
With you here beside me, everything's alright
Can't imagine life alone, don't want to let you go
Stay with me forever, let our love just flow
                    """
                ],
                [
                    "rap piano street tough piercing vocal hip-hop synthesizer clear vocal male",
                    """[verse]
Woke up in the morning, sun is shining bright
Chasing all my dreams, gotta get my mind right
City lights are fading, but my vision's clear
Got my team beside me, no room for fear
Walking through the streets, beats inside my head
Every step I take, closer to the bread
People passing by, they don't understand
Building up my future with my own two hands

[chorus]
This is my life, and I'm aiming for the top
Never gonna quit, no, I'm never gonna stop
Through the highs and lows, I'mma keep it real
Living out my dreams with this mic and a deal
                    """
                ]
            ],
            inputs=[genre_txt, lyrics_txt],
            outputs=[music_out],
            cache_examples=True,
            cache_mode="eager",
            fn=infer
        )

    submit_btn.click(
        fn=infer,
        inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens],
        outputs=[music_out]
    )

demo.queue().launch(show_error=True)