import math
import numpy as np
import os
import re
import tempfile
import torch

from suno_utils.audio import Audio
from suno_utils.audio.conversion import change_audio_speed
from suno_utils.utils.clip import SunoClip

from suno_utils.tasks.mert_25 import encode as encode_semantic, preload_models as preload_semantic_models
from suno_utils.tasks.dac_vae_fixed_25hz import preload_models as preload_codec_models, decode_stream_to_full_audio

from suno_utils.gpt.generation import GenerationConfig, CfgGenerationConfig
from suno_utils.gpt.engine import Engine
from suno_utils.gpt.generation_engine import make_request
from suno_utils.gpt.generation_prompt import ALL_AUDIO_PROMPTS

from suno_utils.diffusion import generation as diffusion_gen
from suno_utils.diffusion.generation import preload_tokenizer
from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job


global engine, diffusion_engine


# gpt_model_fp = "/app/suno/checkpoints/2025-04-28_15-12-21/last_ckpt_infer.pt"  # bct 6b
# gpt_model_fp = "/app/suno/checkpoints/2025-04-23_16-23-43/last_ckpt_infer.pt"  # bct finetune
# gpt_model_fp = "/app/suno/checkpoints/2025-04-26_23-38-19/last_ckpt_infer.pt"  # auk-d (new)
# gpt_model_fp = "/app/suno/checkpoints/2025-04-20_05-36-49/last_ckpt_infer.pt"  # auk-d
# gpt_model_fp = "/app/suno/checkpoints/2025-04-08_19-36-46/last_ckpt_infer.pt"  # (old auk-d)
# gpt_model_fp = "/app/suno/checkpoints/2025-04-05_04-09-06/last_ckpt_infer.pt"  # (old 1-round auk-d)
# gpt_model_fp = "/app/suno/checkpoints/2025-02-04_21-04-31/last_ckpt_infer.pt"  # (6b base)
# gpt_model_fp = "/app/suno/checkpoints/2025-01-28_12-59-39/last_ckpt_infer.pt"  # (3b orig)
# gpt_model_fp = "/app/suno/checkpoints/2024-12-05_22-01-59/last_ckpt_infer.pt"  # (old chef)


def _get_engine():
    global engine
    return engine


def _get_diffusion_engine():
    global diffusion_engine
    return diffusion_engine


def setup_models(gpt_model_fp, diff_model_fp):
    global engine, diffusion_engine
    engine = Engine(
        gpt_model_fp,
        "/app/suno/models/chirp_v2/tokenizer_60k.json",
        max_sequences=8,
        compile=True,
    )
    _ = diffusion_gen.preload_dit_model(
        dit_model_filepath=diff_model_fp,
        use_ema_if_exists=True,
        compile=True,
        weights_precision=torch.bfloat16,
    )
    _ = preload_tokenizer("/app/suno/models/chirp_v2/tokenizer_60k.json")
    _ = preload_semantic_models(
        "/home/georg/notebooks/gpu_nb/tmp/mert_25.pt", 
        "/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy",
    )
    _ = preload_codec_models("/home/georg/notebooks/gpu_nb/tmp/dac_vae_fixed_25hz_2.pth")
    diffusion_engine = UpsampleEngine(min_chunk_size=750)
    # initiate compile
    _ = generate("", "", max_gen_duration_s=9)


def load_audio(audio_fp):
    if re.match(r"^[0-9a-z\-]{36}$", audio_fp):
        a = SunoClip(audio_fp).audio()
    else:
        a = Audio.from_file(audio_fp, sample_rate=48_000, n_channels=2)
    return a


def load_clip(suno_uuid):
    assert re.match(r"^[0-9a-z\-]{36}$", suno_uuid)
    clip = SunoClip(suno_uuid)
    return clip.full_arr()[:,:1], clip.audio().normalize_volume(-14)


def adjust_speed(audio, factor=1.0):
    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_path = os.path.join(tmp_dir, "audio.wav")
        audio.convert(48_000, 2, 2).to_wav(tmp_path)
        out_audio = change_audio_speed(tmp_path, factor, True, 48_000, 2, 2)
    return out_audio


def generate(
    text, 
    tags, 
    gpt_cfg_coef=1.0,
    gpt_cfg_coef_tags=2.5,
    gpt_cfg_coef_max_steps=25*60*2,
    gpt_cfg_coef_tags_max_steps=25*60*2,
    gpt_cfg_coef_audio=0.0,
    diff_n_steps=10,
    control_tags=None,  #"{min_duration:60;max_duration:480}",
    max_gen_duration_s=3*60, 
    n_skip_semantic=1,
    audio_history=None,
    audio_artist=None,
    audio_cover=None,
    audio_vocals=None,
    audio_instrumental=None,
    audio_playlist=None,
    vocals_noise_lvl=0.2,
    gpt_cfg_streams=None,
    gpt_gen_cfg=None, 
    diff_gen_cfg=None,
    return_semantic=False,
    seed=0,
):
    global engine, diffusion_engine
    text = text.strip() if text is None else text.strip()
    tags = tags.strip() if tags is None else tags.strip()
    if len(text.strip()) == 0 and max_gen_duration_s >= 10:
        text = "[instrumental]"
        if "instrumental" not in tags.lower():
            if len(tags) == 0:
                tags = "instrumental"
            else:
                tags = tags + "; instrumental"
    if gpt_gen_cfg is None:
        if audio_history is not None:
            # TODO: diffusion continue not yet implemented
            audio_history = encode_semantic(audio_history)[:,:1]
        if audio_artist is not None:
            audio_artist = encode_semantic(audio_artist.normalize_volume())[:,:1]
        if audio_cover is not None:
            audio_cover = encode_semantic(audio_cover.normalize_volume())[:,:1]
        if audio_vocals is not None:
            if vocals_noise_lvl > 0:
                audio_vocals = audio_vocals.normalize_volume()
                arr = audio_vocals.array_float
                arr = arr + np.random.normal(0, vocals_noise_lvl, arr.shape).astype(arr.dtype)
                audio_vocals = Audio.from_array_float(np.clip(arr, -1.1, 1.1), 48_000)
            audio_vocals = encode_semantic(audio_vocals)[:,:1]
        if audio_instrumental is not None:
            audio_instrumental = encode_semantic(audio_instrumental.normalize_volume())[:,:1]
        if audio_playlist is not None:
            if not isinstance(audio_playlist, list):
                audio_playlist = [audio_playlist]
            tmp_audio_playlist = encode_semantic(audio_playlist[0].normalize_volume())[:,:1]
            for n in range(1, len(audio_playlist)):
                tmp_audio_playlist = np.concatenate([
                    tmp_audio_playlist,
                    np.full((1, 1), engine.model.config.semantic_playlist_token, dtype=tmp_audio_playlist.dtype),
                    encode_semantic(audio_playlist[n].normalize_volume())[:,:1]
                ], axis=0)
            audio_playlist = tmp_audio_playlist
        if gpt_cfg_streams is None:
            gpt_cfg_streams = [
                CfgGenerationConfig(
                    stream_type="tag",
                    prompts=["tag"] + ALL_AUDIO_PROMPTS,
                    null_prompts=ALL_AUDIO_PROMPTS,
                    weight=gpt_cfg_coef_tags,
                    max_steps=gpt_cfg_coef_tags_max_steps,
                ),
            ]
            if gpt_cfg_coef_audio > 0:
                gpt_cfg_streams.append(CfgGenerationConfig(
                    stream_type="custom",
                    prompts=["lyrics", "tag"] + ALL_AUDIO_PROMPTS,
                    null_prompts=["lyrics", "tag"],
                    weight=gpt_cfg_coef_audio,
                    max_steps=None,
                ))
        gpt_gen_cfg = GenerationConfig(
            text=text,
            text_tags=tags,
            history_arr=audio_history,
            artist_arr=audio_artist,
            cover_arr=audio_cover,
            underpaint_arr=audio_vocals,
            overpaint_arr=audio_instrumental,
            playlist_arr=audio_playlist,
            cfg_coef=gpt_cfg_coef,
            cfg_coef_tags=gpt_cfg_coef_tags,
            cfg_coef_max_steps=gpt_cfg_coef_max_steps,
            cfg_coef_tags_max_steps=gpt_cfg_coef_tags_max_steps,
            n_repeat_tags=3 if len(tags) < 50 else 1,
            n_repeat_neg_tags=1,
            n_skip_semantic=n_skip_semantic,
            text_start_control_tags=control_tags,
            cfg_coef_neg_tags=-1.0,
            text_neg_tags="repetitive, loop",
            temp_semantic=0.9,
            top_k_semantic=1500,
            top_p_semantic=None,
            min_p_semantic=0.005,
            max_tag_len=512,
            allow_control_replace_end=False,
            n_batch=1,
            min_eos_p=0.1,
            min_text_offset=0,
            eos_pad_duration_s=0,
            max_gen_duration_s=max_gen_duration_s,
            random_seed=seed,
            custom_null_fields=ALL_AUDIO_PROMPTS,
            cfg_streams=gpt_cfg_streams,
        )

    requests = [
        make_request(f"{i}", gpt_gen_cfg, engine.model.config, engine.tokenizer)
        for i in range(2)
    ]
    jobs = engine.run_request(requests, tqdm_enabled=max_gen_duration_s>=10)
    out_gpt = []
    for n, job in enumerate(jobs):
        stream = engine.token_generator(job)
        arr = torch.stack(list(stream))[:,1]
        if arr[-1] == 4000:
            arr = arr[:-1]
        print(f"{round(arr.shape[-1]/25*n_skip_semantic)}s for track {n}")
        arr2 = torch.zeros(arr.shape[0]*n_skip_semantic, dtype=arr.dtype) + engine.model.config.semantic_pad_token
        arr2[::n_skip_semantic] = arr
        out_gpt.append(arr2)
    if diff_gen_cfg is None:
        diff_gen_cfg = diffusion_gen.DiffusionGenerationConfig(
            lyrics=text,
            tags=tags,
            steps=diff_n_steps,
            text_cfg_coef=2.0,
            codec_scale_factor=0.4,
            scale_ctx_vector=True,
            noise_ctx_level=0.75,
            noise_ctx_pad_len=0,
            drop_semantic_tokens=False,
            seed=24+seed,
        )
    out_audio = []
    for in_sem_arr in out_gpt:
        request = Request(
            id="dummy",
            generation_config=diff_gen_cfg,
            tokens=in_sem_arr,
            input_tokens_finished=True,
        )
        result = diffusion_engine.run_request(request, tqdm_enabled=max_gen_duration_s>=10)
        vae_latents = torch.concat(result.vae_latents)
        a = decode_stream_to_full_audio(vae_latents).normalize_volume(target_db=-14)
        out_audio.append(a)
        if max_gen_duration_s >= 10:
            a.play()
    if return_semantic:
        return out_gpt, out_audio
    return out_audio
        
    
def upsample(
    sem_array_list, 
    text="", 
    tags="", 
    n_steps=16, 
    text_cfg_coef=2.0,
    noise_ctx_level=0.75, 
    ignore_history_every=0, 
    gen_cfg=None, 
    seed=0,
):
    # https://github.com/suno-ai/glockenspiel/pull/9525/files
    assert isinstance(sem_array_list, list)
    if gen_cfg is None:
        gen_cfg = diffusion_gen.DiffusionGenerationConfig(
            lyrics=text,
            tags=tags,
            steps=n_steps,
            text_cfg_coef=text_cfg_coef,
            codec_scale_factor=0.4,
            scale_ctx_vector=True,
            noise_ctx_level=noise_ctx_level,
            noise_ctx_pad_len=0,
            drop_semantic_tokens=False,
            seed=24+seed,
        )
    out_audio = []
    for in_sem_arr in sem_array_list:
        in_kwargs = {
            "id": "dummy", 
            "generation_config": gen_cfg, 
            "tokens": in_sem_arr, 
            "input_tokens_finished": True,
        }
        if ignore_history_every > 0:
            in_kwargs["ignore_history_every"] = ignore_history_every
        request = Request(**in_kwargs)
        result = diffusion_engine.run_request(request, tqdm_enabled=True)
        vae_latents = torch.concat(result.vae_latents)
        a = decode_stream_to_full_audio(vae_latents).normalize_volume(target_db=-14)
        out_audio.append(a)
        a.play()
    return out_audio


def _crossfade(arr_1, arr_2, sr=48000, duration=20):
    assert arr_1.shape[0] == 2 and arr_2.shape[0] == 2  # stereo

    fade_len = sr * duration
    assert arr_1.shape[1] >= fade_len and arr_2.shape[1] >= fade_len, \
        "Each array must be at least as long as the fade duration"

    # Sine-based fade curves
    t = np.linspace(0, np.pi, fade_len)
    fade_out = np.cos(t / 2)
    fade_in = np.sin(t / 2)

    # Crossfade only the overlapping parts
    crossfaded = (
        arr_1[:, -fade_len:] * fade_out +
        arr_2[:, :fade_len] * fade_in
    )

    # Concatenate: arr_1 up to overlap + crossfaded + arr_2 after overlap
    merged = np.concatenate([
        arr_1[:, :-fade_len],
        crossfaded,
        arr_2[:, fade_len:]
    ], axis=1)

    return merged


def upsample_crossfade(
    sem_array_list, 
    text="", 
    tags="", 
    n_steps=16, 
    text_cfg_coef=2.0,
    noise_ctx_level=0.75, 
    window_duration_s=60,
    fade_duration_s=20,
    min_pred_duration_s=30,
    gen_cfg=None, 
    seed=0,
):
    assert isinstance(sem_array_list, list)
    if gen_cfg is None:
        gen_cfg = diffusion_gen.DiffusionGenerationConfig(
            lyrics=text,
            tags=tags,
            steps=n_steps,
            text_cfg_coef=text_cfg_coef,
            codec_scale_factor=0.4,
            scale_ctx_vector=True,
            noise_ctx_level=noise_ctx_level,
            noise_ctx_pad_len=0,
            drop_semantic_tokens=False,
            seed=24+seed,
        )
    window_len = 25 * window_duration_s
    shift_len = 25 * (window_duration_s - fade_duration_s)
    min_window = 25 * min_pred_duration_s
    overlap_len = window_len - shift_len
    out_audio = []
    for in_sem_arr in sem_array_list:
        outs = []
        num_samples = len(in_sem_arr)
        n_chunks = max(1, math.ceil((num_samples - min_window - window_len) / shift_len) + 1)
        for n in range(n_chunks):
            start = n * shift_len
            end = start + window_len if n < n_chunks - 1 else num_samples
            arr_chunk = in_sem_arr[start:end]
            in_kwargs = {
                "id": "dummy", 
                "generation_config": gen_cfg, 
                "tokens": arr_chunk, 
                "input_tokens_finished": True,
            }
            request = Request(**in_kwargs)
            result = diffusion_engine.run_request(request, tqdm_enabled=False)
            vae_latents = torch.concat(result.vae_latents)
            a = decode_stream_to_full_audio(vae_latents)
            outs.append(a)
        out_merged = outs[0].array_float
        for out in outs[1:]:
            out_merged = _crossfade(out_merged, out.array_float)
        a = Audio.from_array_float(out_merged, 48_000).normalize_volume(target_db=-14)
        out_audio.append(a)
        a.play()
    return out_audio


###################
# extra functions #
###################


import math
import random
import re


def _space_repl(m):
    s = m.group()
    n_newline = s.count("\n")
    if n_newline >= 2:
        return "\n\n"
    elif n_newline == 1:
        return "\n"
    return " "


def _simplify_whitespace(text, retain_newlines=True):
    """simplify while respecting up to 2 newlines"""
    if retain_newlines:
        text = re.sub(r"\s+", _space_repl, text).strip()
    else:
        text = re.sub(r"\s+", " ", text).strip()
    return text


def simplify_text(text):
    text = re.sub(r"\[.*?\]", " ", text)
    text = re.sub(r"\s+", " ", text).strip().lower()
    # text = _simplify_whitespace(re.sub(r"\[.*?\]", " ", text), retain_newlines=True)
    return text


def shuffle_lyrics(lyrics, n_grams=1):
    lyrics = lyrics.lower()
    lyrics = re.sub(r"\[.*?\]", " ", lyrics)
    lyrics = re.sub(r"[\s+\,\;\(\)]", " ", lyrics).strip()
    words = lyrics.split()
    fragments = [
        " ".join(words[n*n_grams:(n+1)*n_grams])
        for n in range(math.ceil(len(words)/n_grams))
    ]
    fragments = list(set(fragments))
    random.shuffle(fragments)
    return "; ".join(fragments)

####
# upsample
####


####
# diff source sep
####

STEM_INSTRUMENT_CATEGORIES = [
    "Vocals",
    "Drums",
    "Bass",
    "Guitar",
    "Keyboard",
    "Percussion",
    "Strings",
    "Synth",
    "FX",
    "Brass",
    "Woodwinds",
]

STEM_TYPE_ID_TO_NAME = {
    "0": "Lead Vocal",
    "1": "Electric Guitar",
    "2": "Backing Vocals",
    "3": "Drum Kit",
    "4": "Bass",
    "5": "Piano",
    "6": "Acoustic Guitar",
    "7": "Percussion",
    "8": "String Section",
    "9": "Synth Pad",
    "10": "Synthesizer",
    "11": "Organ",
    "12": "Synth Bass",
    "13": "Lead Electric Guitar",
    "14": "Synth Keys",
    "15": "Rhythm Electric Guitar",
    "16": "Electronic Drum Kit",
    "17": "Noise effects",
    "18": "Electric Piano",
    "19": "Arr. Electric Guitar",
    "20": "Brass section",
    "21": "Distorted Electric Guitar",
    "22": "Upright Bass",
    "23": "Synth Strings",
    "24": "Synth Lead",
    "25": "Rhythm Acoustic Guitar",
    "26": "Intro Count (Click + Key)",
    "27": "Flute",
    "28": "Arpeggiator",
    "29": "Tambourine",
    "30": "Trumpet",
    "31": "Harp",
    "32": "Synth Voice",
    "33": "Accordion",
    "34": "Fiddle",
    "35": "Synth Brass",
    "36": "Violin",
    "37": "Pedal Steel Guitar",
    "38": "Sound effects",
    "39": "Brass Instruments",
    "40": "Digital Piano",
    "41": "Tenor Saxophone",
    "42": "Mandolin",
    "43": "Trombone",
    "44": "Clarinet",
    "45": "French Horn",
    "46": "Banjo",
    "47": "Lead Acoustic Guitar",
    "48": "Rhythm Electric Guitar (Arpeggio)",
    "49": "Glockenspiel",
    "50": "Shaker",
    "51": "Electric Bass",
    "52": "Timpani",
    "53": "Vibes",
    "54": "Harmonica",
    "55": "Hand Clap",
    "56": "Guitar Synth",
    "57": "Lap Steel Guitar",
    "58": "Arr. Acoustic Guitar",
    "59": "Cello",
    "60": "Electronic Percussion",
    "61": "Woodwinds",
    "62": "Oboe",
    "63": "Rhythm Acoustic Guitar (Arpeggio)",
    "64": "Saxophone",
    "65": "Drums and Percussion",
    "66": "Female Backing Vocals",
    "67": "Marimba",
    "68": "Male Backing Vocals",
    "69": "Alto Saxophone",
    "70": "Bells",
    "71": "Pitched Percussion",
    "72": "Synth Ambiant",
    "73": "Dobro",
    "74": "Celesta",
    "75": "Congas",
    "76": "Synth Flute",
    "77": "Harpsichord",
    "78": "Baritone Saxophone",
    "79": "Orchestra Hit",
    "80": "Female Lead Vocal",
    "81": "Male Lead Vocal",
    "82": "Keyboard",
    "83": "Wind Chimes",
    "84": "Wind Instruments",
    "85": "Double Bass",
    "86": "Acoustic Drum Kit",
    "87": "Orchestral Percussion",
    "88": "Xylophone",
    "89": "Ukulele",
    "90": "Orchestra",
    "91": "Sound Effects",
    "92": "Cymbals",
    "93": "Tuba",
    "94": "Clap",
    "95": "Latin Percussion",
    "96": "Intro Count (Click)",
    "97": "Bassoon",
    "98": "Triangle",
    "99": "Snap (Fingers)",
    "100": "Sitar",
    "101": "Tubular Bell",
    "102": "Bongos",
    "103": "Cowbell",
    "104": "Sample",
    "105": "Metal bars",
    "106": "Steel Drums",
    "107": "Guitar",
    "108": "Viola",
    "109": "Soprano Saxophone",
    "110": "Claves/Woodblock",
    "111": "Whistle",
    "112": "English Horn",
    "113": "Jingle Bells",
    "114": "Thumb Piano",
    "115": "Piccolo",
    "116": "Intro Count (Key)",
    "117": "Bagpipes",
    "118": "Voice",
    "119": "Bass Drum",
    "120": "Music Box",
    "121": "Rhodes",
    "122": "Appalachian Dulcimer",
    "123": "Human Beatbox",
    "124": "Flugelhorn",
    "125": "Fretless Bass",
    "126": "Theremin",
    "127": "Vocoder",
}

import os
import modal
import json
import tempfile
from suno_utils.audio import Audio
from suno_utils.utils.s3 import _upload_s3_file
from uuid import uuid4
from contextlib import redirect_stderr, redirect_stdout
import requests

def split_stems(fp):
    uid = str(uuid4())
    s3_filepath = f"s3://suno-data/georg/tmp/modal/{uid}.mp3"
    with tempfile.TemporaryDirectory() as tmp_dir:
        local_filepath = os.path.join(tmp_dir, f"{uid}.mp3")
        _ = Audio.from_file(fp, sample_rate=48_000, n_channels=2).to_hq_mp3(local_filepath)
        _upload_s3_file(local_filepath, s3_filepath)
    in_info = {
        "id": uid,
        "prompt_audio": s3_filepath,
        "model_name": "stems",
        "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "stem_type_group_name": "Vocals",
#             "stem_type_id": 0,
            "multi_ids": [f"{uid}_stem", f"{uid}_complement"],
        },
    }
    with redirect_stderr(open(os.devnull, "w")):
        with redirect_stdout(open(os.devnull, "w")):
            model_f_stem = modal.Cls.lookup("stems-stems_v1-dev", "StemStub")
            model_f_stem.stem.remote(json.dumps(in_info))
    out_stem_fp = f"https://cdn1.suno.ai/{uid}_stem.opus"
    out_complement_fp = f"https://cdn1.suno.ai/{uid}_complement.opus"
    r_stem = requests.get(out_stem_fp)
    r_complement = requests.get(out_complement_fp)
    with tempfile.TemporaryDirectory() as tmp_dir:
        out_local_stem_fp = os.path.join(tmp_dir, f"{uid}_stem.opus")
        out_local_complement_fp = os.path.join(tmp_dir, f"{uid}_complement.opus")
        with open(out_local_stem_fp, "wb") as f:
            f.write(r_stem.content)
        with open(out_local_complement_fp, "wb") as f:
            f.write(r_complement.content)
        audio_stem = Audio.read_opus(out_local_stem_fp)
        audio_complement = Audio.read_opus(out_local_complement_fp)
    return audio_stem, audio_complement