import tempfile
import torchaudio
import boto3
import json
from dataclasses import dataclass
from typing import Optional, List
import numpy as np
import torch
import os
import statistics
import librosa

from suno_utils.gpt import chirp_v2_5 as chirp_v3
from suno_utils.diffusion import generation as diffusion_gen
from suno_utils.gpt.engine import Engine
from suno_utils.gpt.generation_engine import make_request, semantic_codes

from suno_utils.tasks.upsample_engine import (
    UpsampleEngine,
    Request,
    DiffusionGenerationConfig,
)
from suno_utils.audio import Audio
from suno_utils.worker.loader import _get_compatible_tokens

from suno_utils.tasks import ss_vad
from suno_utils.models.ditto_v2.ditto_v2 import Ditto
from suno_utils.tasks.dac_vae_100hz_peaq import preload_models as preload_vae_models
from suno_utils.gpt.generation import CfgGenerationConfig

import sys

sys.path.insert(0, "/home/sara/neon/sunoDiff/")
from generation import preload_models as preload_diff_models, generate

diff_model_fp = "/app/suno/data/dpo/models/diff_vae_25_peaq_v4_jan28.pt"

USE_COMPILE = False
MAX_STREAMS = 20
N_BATCH = 10
PARALLEL = 4
MIN_CHUNK_SIZE = 25 * 30
MAX_DURATION = 240
SILENCE_CUTOFF = 30
SEP_SAMPLE_RATE = 44100
DITTO_SR = 24000
ENCODER_RATE = 25
DITTO_EMBEDDING_DIM = 128

RUN_COVER = True
RUN_ARTIST = False


# t25 -- better online on artist +3, 0 on cover:  /app/suno/checkpoints/2025-01-14_05-33-54
# t27 -- better online on cover +2, 0 on artist: /app/suno/checkpoints/2025-01-19_05-27-38

DATA_PATH = "/home/sara/glockenspiel/suno_utils/task_eval/test_task.json"
MODEL_NAME = "model_45"  # "model_30b_t5_v1"
GPT_CKPT_PATH = "/app/suno/checkpoints/2025-01-28_12-59-39/last_ckpt_infer.pt"  # f"s3://suno-data/tony/tmp/{MODEL_NAME}.pt"  # 30b_t6

DIFFUSION_CKPT_PATH = (
    "s3://suno-data/tony/tmp/diff/dit_v2_dpo_t2_v1_3k.pt"  # 30b_t6 diffusion default
)
DITTO_S3_PATH = "s3://suno-data/minz/models/ditto_v2_epoch_57.pt"
VAE_S3_PATH = "s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth"


@dataclass
class Prompt:
    lyrics: str
    tags: str
    s3_id: str
    test_tags: Optional[str] = None
    test_lyrics: Optional[str] = None
    neg_tags: Optional[str] = None
    cover_arr: Optional[np.ndarray] = None
    artist_arr: Optional[np.ndarray] = None


def load_prompt_arr(audio_id):
    with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file:
        target_version = "4.0"
        tokens = _get_compatible_tokens(temp_file.name, audio_id, target_version)
    return tokens[: (ENCODER_RATE * MAX_DURATION)]


def run_gpt(prompt: Prompt):
    n_skip_semantic = 1
    gen_config = chirp_v3.GenerationConfig(
        text=prompt.test_lyrics or prompt.lyrics,
        text_tags=prompt.test_tags or prompt.tags,
        text_neg_tags=prompt.neg_tags,
        cfg_coef=1.2,
        # cfg_coef_max_steps=25 * 30,
        artist_arr=prompt.artist_arr,
        cover_arr=prompt.cover_arr,
        n_batch=1,
        n_repeat_tags=3,
        n_skip_semantic=n_skip_semantic,
        min_text_offset=0,
        eos_pad_duration_s=0,
        max_gen_duration_s=MAX_DURATION,
        cfg_streams=[
            CfgGenerationConfig(stream_type="tag", weight=2.0, max_steps=25 * 30),
            CfgGenerationConfig(stream_type="neg_tag", weight=-1.0, max_steps=25 * 30),
            # CfgGenerationConfig(stream_type="artist", weight=2.0, max_steps=25 * 30),
        ],
    )

    requests = [
        make_request(f"{i}", gen_config, gpt_engine.model.config, gpt_engine.tokenizer)
        for i in range(N_BATCH)
    ]

    gpt_outputs = []
    for idx in range(0, len(requests), PARALLEL):
        inner_requests = requests[idx : (idx + PARALLEL)]
        for n, job in enumerate(
            gpt_engine.run_request(inner_requests, tqdm_enabled=True)
        ):
            stream = gpt_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}")
            # do stuff incase skip
            arr2 = (
                torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype)
                + cfg.semantic_pad_token
            )
            arr2[::n_skip_semantic] = arr
            gpt_outputs.append(arr2)

    return gpt_outputs


def run_diffusion(prompt: Prompt, gpt_outputs: List):
    diff_outputs = []
    for in_sem_arr in gpt_outputs:
        audio = generate(
            in_sem_arr,
            lyrics=prompt.test_lyrics or prompt.lyrics,
            tags=prompt.test_tags or prompt.tags,
            text_cfg_coef=1.0,
            steps=16,
            seed=0,
        )
        diff_outputs.append(audio)

    return diff_outputs


def load_audio(audio_path, duration=MAX_DURATION):
    s3 = boto3.client("s3")

    if os.path.exists(audio_path):
        waveform, sr = torchaudio.load(audio_path)
    else:
        with tempfile.NamedTemporaryFile(suffix=".mp3") as temp_file:
            s3.download_file(
                "suno-data-uploads", f"studio/uploads/{audio_path}.mp3", temp_file.name
            )
            waveform, sr = torchaudio.load(temp_file.name)

    waveform = torch.mean(waveform, dim=0).unsqueeze(0)[:, : sr * duration]
    if sr != SEP_SAMPLE_RATE:
        resampler = torchaudio.transforms.Resample(sr, SEP_SAMPLE_RATE)
        waveform = resampler(waveform)

    return waveform


def load_audio_from_gen(audio, duration=MAX_DURATION, trim=True):
    with tempfile.NamedTemporaryFile(suffix=".mp3") as temp_file:
        audio.get_segment(to_s=MAX_DURATION).write_hq_mp3(temp_file.name)
        waveform = load_audio(temp_file.name, duration=duration)
    return waveform


def cosine_similarity(a, b):
    # Calculate dot product
    dot_product = np.dot(a, b)

    # Calculate magnitudes
    magnitude_a = np.sqrt(np.dot(a, a))
    magnitude_b = np.sqrt(np.dot(b, b))

    # Calculate cosine similarity
    return dot_product / (magnitude_a * magnitude_b)


def load_vocals(audio_path, duration=MAX_DURATION, trim=True):
    waveform = load_audio(audio_path, duration=duration)
    vocals = ss_vad.encode(waveform)
    if trim:
        vocals, _ = librosa.effects.trim(vocals, top_db=SILENCE_CUTOFF)
    vocals = torch.from_numpy(vocals)
    return vocals


def load_vocals_from_gen(audio, duration=MAX_DURATION, trim=True):
    with tempfile.NamedTemporaryFile(suffix=".mp3") as temp_file:
        audio.get_segment(to_s=MAX_DURATION).write_hq_mp3(temp_file.name)
        waveform = load_audio(temp_file.name, duration=duration)
    vocals = ss_vad.encode(waveform)
    if trim:
        vocals, _ = librosa.effects.trim(vocals, top_db=SILENCE_CUTOFF)
    vocals = torch.from_numpy(vocals)
    return vocals


def embed_waveform(waveform, task):
    resampler = torchaudio.transforms.Resample(SEP_SAMPLE_RATE, DITTO_SR)
    return (
        ditto.music_to_latent(resampler(waveform), task=task)[0].detach().cpu().numpy()
    )


if __name__ == "__main__":
    gpt_engine = Engine(
        GPT_CKPT_PATH,
        #     "/app/suno/checkpoints/2025-01-26_01-35-19/last_ckpt_infer.pt",  # skip
        # "/app/suno/checkpoints/2025-01-28_12-59-39/last_ckpt_infer.pt",  # base long
        #     "/app/suno/checkpoints/2025-01-18_18-03-13/last_ckpt_infer.pt",  # base
        #     "/app/suno/checkpoints/2024-12-05_22-01-59/last_ckpt_infer.pt",  # (old chef)
        "/app/suno/models/chirp_v2/tokenizer_60k.json",
        max_sequences=MAX_STREAMS,
        compile=False,
        max_length_s=240,
    )
    cfg = gpt_engine.model.config

    _ = preload_diff_models(
        tokenizer_filepath="/home/georg/notebooks/gpu_nb/tmp/tokenizer_60k.json",
        semantic_model_filepath="/home/georg/notebooks/gpu_nb/tmp/mert_25.pt",
        semantic_clusters_filepath="/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy",
        codec_filepath="/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth",
        dit_model_filepath=diff_model_fp,
        model_type="prefix",
        weights_precision=torch.bfloat16,
        compile=True,
    )

    ditto_path = chirp_v3._get_model_if_needed(DITTO_S3_PATH)
    ditto = Ditto(
        latent_dim=DITTO_EMBEDDING_DIM,
        model_path=ditto_path,
        is_flash=False,
        is_serving=True,
    )

    ditto = ditto.eval().to("cuda:1")

    # ss_vad_config_path = chirp_v3._get_model_if_needed(ss_vad.YAML_PATH)
    # ss_vad_model_path = chirp_v3._get_model_if_needed(ss_vad.MODEL_PATH)
    # ss_vad.preload_models(
    #     checkpoint_filepath=ss_vad_model_path,
    #     config_path=ss_vad_config_path,
    # )

    # preload_vae_models(VAE_S3_PATH)
    print("Finish loading models")

    with open(DATA_PATH, "r") as file:
        data = json.load(file)

    test_covers = data["cover"]
    test_artist = data["artist"]

    if RUN_ARTIST:
        TASK = "artist_vox_sim"
        TRIM = True

        scores = {}
        for artist in test_artist:
            prompt = Prompt(**artist)
            prompt.artist_arr = load_prompt_arr(prompt.s3_id)

            gpt_outputs = run_gpt(prompt)
            diff_outputs = run_diffusion(prompt, gpt_outputs)

            S3_ID = prompt.s3_id
            source_wav = load_vocals(S3_ID)
            source_embed = embed_waveform(source_wav, TASK)
            scores[S3_ID] = []
            for full_audio in diff_outputs:
                artist_wav = load_vocals_from_gen(full_audio)
                artist_embed = embed_waveform(artist_wav, TASK)
                score = cosine_similarity(source_embed, artist_embed)
                scores[S3_ID].append(score)

        for key in scores:
            avg_score = statistics.mean(scores[key])
            min_score = min(scores[key])
            max_score = max(scores[key])
            print(f"{key}: average: {avg_score} min: {min_score} max: {max_score}")

        np.savez(
            f"{MODEL_NAME}_{TASK}_artist_{MAX_DURATION}_{SILENCE_CUTOFF}.npz",
            **scores,
        )

    if RUN_COVER:
        TASK = "self_sim"

        scores = {}
        for cover in test_covers:
            prompt = Prompt(**cover)
            prompt.cover_arr = load_prompt_arr(prompt.s3_id)
            print(f"Generating {prompt.s3_id} with {prompt.test_tags}")

            gpt_outputs = run_gpt(prompt)
            diff_outputs = run_diffusion(prompt, gpt_outputs)

            S3_ID = prompt.s3_id
            source_wav = load_audio(S3_ID)
            source_embed = embed_waveform(source_wav, TASK)
            scores[S3_ID] = []
            for full_audio in diff_outputs:
                cover_wav = load_audio_from_gen(full_audio)
                cover_embed = embed_waveform(cover_wav, TASK)
                score = cosine_similarity(source_embed, cover_embed)
                scores[S3_ID].append(score)

        for key in scores:
            avg_score = statistics.mean(scores[key])
            min_score = min(scores[key])
            max_score = max(scores[key])
            print(f"{key}: average: {avg_score} min: {min_score} max: {max_score}")

        np.savez(f"{MODEL_NAME}_{TASK}_cover_{MAX_DURATION}.npz", **scores)