import gradio as gr
import subprocess
import os
import shutil
import tempfile
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
import torch

is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False

# Install required package
def install_flash_attn():
    try:
        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,
        )
        print("flash-attn installed successfully!")
    except subprocess.CalledProcessError as e:
        print(f"Failed to install flash-attn: {e}")
        exit(1)

# Install flash-attn
install_flash_attn()

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"
)

# Add xcodec_mini_infer and descriptaudiocodec to sys path
import sys
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'))


import os
import sys
import torch
import numpy as np
import json
import re
import uuid
import gradio as gr
from tqdm import tqdm
from omegaconf import OmegaConf
import torchaudio
from torchaudio.transforms import Resample
import soundfile as sf
from einops import rearrange
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

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'))
from codecmanipulator import CodecManipulator
from mmtokenizer import _MMSentencePieceTokenizer

# Load models once at startup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

# Load tokenizers and codec tools
print("Loading tokenizers...")
mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
codectool = CodecManipulator("xcodec", 0, 1)

# Load codec models
print("Loading codec models...")
model_config = OmegaConf.load('./xcodec_mini_infer/final_ckpt/config.yaml')
codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
parameter_dict = torch.load('./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', map_location='cpu')
codec_model.load_state_dict(parameter_dict['codec_model'])
codec_model.to(device)
codec_model.eval()

# Load vocoders
print("Loading vocoders...")
vocal_decoder, inst_decoder = build_codec_model(
    './xcodec_mini_infer/decoders/config.yaml',
    './xcodec_mini_infer/decoders/decoder_131000.pth',
    './xcodec_mini_infer/decoders/decoder_151000.pth'
)

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

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]

def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    limit = 0.99
    max_val = wav.abs().max()
    wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
    torchaudio.save(path, wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)

@spaces.GPU(duration=150)
def run_inference(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=2000):
    try:
        # Create temporary output directory
        output_dir = tempfile.mkdtemp()
        stage1_output_dir = os.path.join(output_dir, "stage1")
        os.makedirs(stage1_output_dir, exist_ok=True)

        # Process inputs
        structured_lyrics = split_lyrics(lyrics_txt_content)
        full_lyrics = "\n".join(structured_lyrics)
        prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genre_txt_content}\n{full_lyrics}"] + structured_lyrics

        # Generation parameters
        top_p = 0.93
        temperature = 1.0
        repetition_penalty = 1.2
        start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
        end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
        run_n_segments = min(num_segments + 1, len(structured_lyrics))

        # Generate tokens
        raw_output = None
        for i in tqdm(range(1, run_n_segments)):
            section_text = prompt_texts[i].replace('[start_of_segment]', '').replace('[end_of_segment]', '')
            guidance_scale = 1.5 if i <= 1 else 1.2
            prompt_ids = start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
            prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)
            
            input_ids = prompt_ids if i == 1 else torch.cat([raw_output, prompt_ids], dim=1)
            if input_ids.shape[-1] > 16384 - max_new_tokens - 1:
                input_ids = input_ids[:, -(16384 - max_new_tokens - 1):]

            with torch.no_grad():
                output_seq = model.generate(
                    input_ids=input_ids,
                    max_new_tokens=max_new_tokens,
                    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,
                )

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

        # Process generated tokens
        ids = raw_output[0].cpu().numpy()
        soa_idx = np.where(ids == mmtokenizer.soa)[0]
        eoa_idx = np.where(ids == mmtokenizer.eoa)[0]
        vocals, instrumentals = [], []

        for i in range(len(soa_idx)):
            codec_ids = ids[soa_idx[i]+1:eoa_idx[i]]
            codec_ids = codec_ids[:2 * (len(codec_ids) // 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]))

        # Generate audio
        vocals = np.concatenate(vocals, axis=1)
        instrumentals = np.concatenate(instrumentals, axis=1)
        
        with torch.no_grad():
            vocal_audio = codec_model.decode(torch.tensor(vocals.astype(np.int16)).long().unsqueeze(0).permute(1, 0, 2).to(device))
            inst_audio = codec_model.decode(torch.tensor(instrumentals.astype(np.int16)).long().unsqueeze(0).permute(1, 0, 2).to(device))

        # Mix and save audio
        final_audio = (vocal_audio.cpu().squeeze() + inst_audio.cpu().squeeze()) / 2
        output_path = os.path.join(output_dir, "final_output.wav")
        save_audio(final_audio.unsqueeze(0), output_path, 16000)

        return output_path

    except Exception as e:
        print(f"Error during inference: {str(e)}")
        raise gr.Error(f"Generation failed: {str(e)}")

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# YuE Music Generator")
    with gr.Row():
        with gr.Column():
            genre_txt = gr.Textbox(label="Genre Tags", placeholder="e.g., female vocal, jazz, piano")
            lyrics_txt = gr.Textbox(label="Lyrics", lines=10, placeholder="Enter lyrics with sections like [verse], [chorus]")
            num_segments = gr.Slider(1, 10, value=2, label="Number of Segments")
            max_tokens = gr.Slider(500, 3000, value=2000, label="Max Tokens")
            btn = gr.Button("Generate Music")
        with gr.Column():
            audio_out = gr.Audio(label="Generated Music")

    examples = 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"""]
        ],
        inputs=[genre_txt, lyrics_txt],
        outputs=audio_out
    )

    btn.click(
        fn=run_inference,
        inputs=[genre_txt, lyrics_txt, num_segments, max_tokens],
        outputs=audio_out
    )

if __name__ == "__main__":
    demo.launch()