import gradio as gr
import subprocess
import os
import shutil
import tempfile
import spaces
import torch
import torch.nn.functional as F
import sys

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

import uuid
from tqdm import tqdm
from einops import rearrange
from codecmanipulator import CodecManipulator
from mmtokenizer import _MMSentencePieceTokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
import glob
import time
import copy
from collections import Counter
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
import re
import multiprocessing

def empty_output_folder(output_dir):
    # List all files in the output directory
    files = os.listdir(output_dir)

    # Iterate over the files and remove them
    for file in files:
        file_path = os.path.join(output_dir, file)
        try:
            if os.path.isdir(file_path):
                # If it's a directory, remove it recursively
                shutil.rmtree(file_path)
            else:
                # If it's a file, delete it
                os.remove(file_path)
        except Exception as e:
            print(f"Error deleting file {file_path}: {e}")

device = "cuda:0"

# --- Model Loading and Quantization ---
model = AutoModelForCausalLM.from_pretrained(
    "m-a-p/YuE-s1-7B-anneal-en-cot",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",  # To enable flashattn, you have to install flash-attn
)
model.to(device)
model.eval()

# Apply dynamic quantization
model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

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'

mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")

codectool = CodecManipulator("xcodec", 0, 1)
model_config = OmegaConf.load(basic_model_config)
codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
parameter_dict = torch.load(resume_path, map_location='cpu')
codec_model.load_state_dict(parameter_dict['codec_model'])
codec_model.to(device)
codec_model.eval()

# --- Parallel Audio Processing ---
def process_audio_wrapper(args):
    # Unpack arguments and call the original process_audio function
    npy, output_path, rescale, other_args, decoder, codec_model = args
    return process_audio(npy, output_path, rescale, other_args, decoder, codec_model)

def parallel_process_audio(stage1_output_set, vocoder_stems_dir, rescale, other_args, vocal_decoder, inst_decoder,
                           codec_model, num_processes=4):
    with multiprocessing.Pool(processes=num_processes) as pool:
        tasks = []
        for npy in stage1_output_set:
            if 'instrumental' in npy:
                output_path = os.path.join(vocoder_stems_dir, 'instrumental.mp3')
                decoder = inst_decoder
            else:
                output_path = os.path.join(vocoder_stems_dir, 'vocal.mp3')
                decoder = vocal_decoder
            tasks.append((npy, output_path, rescale, other_args, decoder, codec_model))

        results = pool.map(process_audio_wrapper, tasks)

    return results

# --- Optimized Music Generation ---
def generate_music(
        max_new_tokens=5,
        run_n_segments=2,
        genre_txt=None,
        lyrics_txt=None,
        use_audio_prompt=False,
        audio_prompt_path="",
        prompt_start_time=0.0,
        prompt_end_time=30.0,
        output_dir="./output",
        rescale=False,
        beam_width=3,  # Add beam search
        length_penalty=1.0,  # Add length penalty
        repetition_penalty=1.5, # Add repetition penalty
        batch_size=2
):
    if use_audio_prompt and not audio_prompt_path:
        raise FileNotFoundError(
            "Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
    max_new_tokens = max_new_tokens * 100
    stage1_output_dir = os.path.join(output_dir, f"stage1")
    os.makedirs(stage1_output_dir, exist_ok=True)

    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 load_audio_mono(filepath, sampling_rate=16000):
        audio, sr = torchaudio.load(filepath)
        # Convert to mono
        audio = torch.mean(audio, dim=0, keepdim=True)
        # Resample if needed
        if sr != sampling_rate:
            resampler = Resample(orig_freq=sr, new_freq=sampling_rate)
            audio = resampler(audio)
        return audio

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

    # Call the function and print the result
    stage1_output_set = []

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

    random_id = uuid.uuid4()
    output_seq = None
    # Here is suggested decoding config
    top_p = 0.93
    temperature = 1.0
    # special tokens
    start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
    end_of_segment = mmtokenizer.tokenize('[end_of_segment]')

    raw_output = None
    segment_cache = {}  # Cache for repeated segments

    # Format text prompt
    run_n_segments = min(run_n_segments + 1, len(lyrics))

    print(list(enumerate(tqdm(prompt_texts[:run_n_segments]))))

    # Modified loop for batching and caching
    for i in range(1, run_n_segments, batch_size):
        batch_segments = []
        batch_prompts = []
        for j in range(i, min(i + batch_size, run_n_segments)):
            section_text = prompt_texts[j].replace('[start_of_segment]', '').replace('[end_of_segment]', '')

            # Check cache
            if section_text in segment_cache:
                cached_output = segment_cache[section_text]
                if j > 1:
                    raw_output = torch.cat([raw_output, cached_output], dim=1)
                else:
                    raw_output = cached_output
                continue

            batch_segments.append(section_text)
            guidance_scale = 1.5 if j <= 1 else 1.2

            if j == 1:
                if use_audio_prompt:
                    audio_prompt = load_audio_mono(audio_prompt_path)
                    audio_prompt.unsqueeze_(0)
                    with torch.no_grad():
                        raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5)
                    raw_codes = raw_codes.transpose(0, 1)
                    raw_codes = raw_codes.cpu().numpy().astype(np.int16)
                    # Format audio prompt
                    code_ids = codectool.npy2ids(raw_codes[0])
                    audio_prompt_codec = code_ids[
                                         int(prompt_start_time * 50): int(prompt_end_time * 50)]  # 50 is tps of xcodec
                    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 + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
            else:
                prompt_ids = end_of_segment + 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 = torch.cat([raw_output, prompt_ids], dim=1) if j > 1 else prompt_ids

            # Use window slicing in case output sequence exceeds the context of model
            max_context = 16384 - max_new_tokens - 1
            if input_ids.shape[-1] > max_context:
                print(
                    f'Section {j}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
                input_ids = input_ids[:, -(max_context):]

            batch_prompts.append(input_ids)

        if not batch_prompts:
            continue  # All segments in the batch were cached

        # Pad prompts in the batch to the same length
        max_len = max(p.size(1) for p in batch_prompts)
        padded_prompts = []
        for p in batch_prompts:
            pad_len = max_len - p.size(1)
            padded_prompt = F.pad(p, (0, pad_len), value=mmtokenizer.eoa)
            padded_prompts.append(padded_prompt)

        batch_input_ids = torch.cat(padded_prompts, dim=0)

        with torch.no_grad():
            output_seqs = model.generate(
                input_ids=batch_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,
                num_beams=beam_width,  # Use beam search
                length_penalty=length_penalty,  # Apply length penalty
            )

        # Process each output in the batch
        for k, output_seq in enumerate(output_seqs):
            if output_seq[0][-1].item() != mmtokenizer.eoa:
                tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
                output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
            if i > 1:
                raw_output = torch.cat([raw_output, batch_prompts[k][:, :batch_input_ids.shape[-1]],
                                        output_seq[:, batch_input_ids.shape[-1]:]], dim=1)
            else:
                raw_output = output_seq

            # Cache the generated output if not already cached
            if batch_segments[k] not in segment_cache:
                segment_cache[batch_segments[k]] = output_seq[:, batch_input_ids.shape[-1]:].cpu()

    # save raw output and check sanity
    ids = raw_output[0].cpu().numpy()
    soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
    eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
    if len(soa_idx) != len(eoa_idx):
        raise ValueError(f'invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')

    vocals = []
    instrumentals = []
    range_begin = 1 if use_audio_prompt else 0
    for i in range(range_begin, 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_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0])
        vocals.append(vocals_ids)
        instrumentals_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1])
        instrumentals.append(instrumentals_ids)
    vocals = np.concatenate(vocals, axis=1)
    instrumentals = np.concatenate(instrumentals, axis=1)
    vocal_save_path = os.path.join(stage1_output_dir,
                                   f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_vocal_{random_id}".replace(
                                       '.', '@') + '.npy')
    inst_save_path = os.path.join(stage1_output_dir,
                                  f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_instrumental_{random_id}".replace(
                                      '.', '@') + '.npy')
    np.save(vocal_save_path, vocals)
    np.save(inst_save_path, instrumentals)
    stage1_output_set.append(vocal_save_path)
    stage1_output_set.append(inst_save_path)

    print("Converting to Audio...")

    # convert audio tokens to audio
    def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
        folder_path = os.path.dirname(path)
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)
        limit = 0.99
        max_val = wav.abs().max()
        wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
        torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)

    # reconstruct tracks
    recons_output_dir = os.path.join(output_dir, "recons")
    recons_mix_dir = os.path.join(recons_output_dir, 'mix')
    os.makedirs(recons_mix_dir, exist_ok=True)
    tracks = []
    for npy in stage1_output_set:
        codec_result = np.load(npy)
        decodec_rlt = []
        with torch.no_grad():
            decoded_waveform = codec_model.decode(
                torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(
                    device))
        decoded_waveform = decoded_waveform.cpu().squeeze(0)
        decodec_rlt.append(torch.as_tensor(decoded_waveform))
        decodec_rlt = torch.cat(decodec_rlt, dim=-1)
        save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy))[0] + ".mp3")
        tracks.append(save_path)
        save_audio(decodec_rlt, save_path, 16000)
    # mix tracks
    for inst_path in tracks:
        try:
            if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
                    and 'instrumental' in inst_path:
                # find pair
                vocal_path = inst_path.replace('instrumental', 'vocal')
                if not os.path.exists(vocal_path):
                    continue
                # mix
                recons_mix = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace('instrumental', 'mixed'))
                vocal_stem, sr = sf.read(inst_path)
                instrumental_stem, _ = sf.read(vocal_path)
                mix_stem = (vocal_stem + instrumental_stem) / 1
                sf.write(recons_mix, mix_stem, sr)
        except Exception as e:
            print(e)

    # vocoder to upsample audios
    vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
    vocoder_output_dir = os.path.join(output_dir, 'vocoder')
    vocoder_stems_dir = os.path.join(vocoder_output_dir, 'stems')
    vocoder_mix_dir = os.path.join(vocoder_output_dir, 'mix')
    os.makedirs(vocoder_mix_dir, exist_ok=True)
    os.makedirs(vocoder_stems_dir, exist_ok=True)

    # Use parallel processing for vocoding
    parallel_process_audio(stage1_output_set, vocoder_stems_dir, rescale, argparse.Namespace(**locals()), vocal_decoder,
                           inst_decoder, codec_model)

    # mix tracks after parallel processing
    instrumental_output_path = os.path.join(vocoder_stems_dir, 'instrumental.mp3')
    vocal_output_path = os.path.join(vocoder_stems_dir, 'vocal.mp3')

    if os.path.exists(instrumental_output_path) and os.path.exists(vocal_output_path):
      instrumental_output, sr = torchaudio.load(instrumental_output_path)
      vocal_output, _ = torchaudio.load(vocal_output_path)
      try:
          mix_output = instrumental_output + vocal_output
          vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
          save_audio(mix_output, vocoder_mix, 44100, rescale)
          print(f"Created mix: {vocoder_mix}")
      except RuntimeError as e:
          print(e)
          print(f"mix {vocoder_mix} failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
    else:
      print("Skipping mix creation, instrumental or vocal output missing.")

    # Post process
    replace_low_freq_with_energy_matched(
        a_file=recons_mix,  # 16kHz
        b_file=vocoder_mix,  # 48kHz
        c_file=os.path.join(output_dir, os.path.basename(recons_mix)),
        cutoff_freq=5500.0
    )
    print("All process Done")
    return recons_mix

@spaces.GPU(duration=120)
def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=5):
    # Ensure the output folder exists
    output_dir = "./output"
    os.makedirs(output_dir, exist_ok=True)
    print(f"Output folder ensured at: {output_dir}")

    empty_output_folder(output_dir)

    # Execute the command
    try:
        music = generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments,
                               output_dir=output_dir, cuda_idx=0, max_new_tokens=max_new_tokens)
        return music
    except Exception as e:
        gr.Warning("An Error Occured: " + str(e))
        return None
    finally:
        print("Temporary files deleted.")

# Gradio

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