import os

import torch
import numpy as np
from suno_utils.audio import Audio
from suno_utils.gpt.chirp_v2_5 import (
    _get_model_if_needed,
)
from suno_utils.gpt.engine import Engine
from suno_utils.tasks.dac_2c_12cb import (
    preload_models as preload_codec_models,
    encode as codec_encode,
)
from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as semantic_encode,
)

from suno_utils.tasks.upsample_engine import (
    UpsampleEngine,
    Request,
    DiffusionGenerationConfig,
)
from suno_utils.diffusion import generation as diffusion_gen

os.environ["CUDA_VISIBLE_DEVICES"] = "5"

torch._logging.set_logs(recompiles=True, graph_breaks=True)  # , guards=True)


def _interleave(semantic_arr, n_factor=1):
    new_semantic_arr = (
        np.zeros(
            (semantic_arr.shape[0] * n_factor, semantic_arr.shape[-1]),
            dtype=semantic_arr.dtype,
        )
        + cfg.semantic_vocab_size
        - 1
    )
    new_semantic_arr[::n_factor] = semantic_arr
    return new_semantic_arr


def process_audio(audio, cfg, n_factor=1):
    audio = audio.normalize_volume(-16)
    sem_arr = semantic_encode(audio, device="cpu")
    if n_factor > 1:
        sem_arr = _interleave(sem_arr, n_factor=n_factor)
    coarse_arr = codec_encode(audio)
    n_frames = min(sem_arr.shape[0], coarse_arr.shape[0])
    sem_arr = sem_arr[:n_frames, : cfg.semantic_n_codebooks]
    coarse_arr = coarse_arr[:n_frames, : cfg.coarse_n_codebooks]

    a_arr = np.concatenate([sem_arr, coarse_arr], axis=-1)
    return a_arr


def load_audio(fp):
    return Audio.from_file(fp, n_channels=2, sample_rate=48_000, byte_width=2)


N_BATCH = 2
MAX_STREAMS = N_BATCH * 4

# preload codec
_ = preload_codec_models("/app/suno/models/chirp_v2/dac_2c_25x12.pt")

# preload mert
_ = preload_semantic_models(
    checkpoint_filepath="/app/suno/models/chirp_v2/mert_25.pt",
    centroids_filepath="/app/suno/models/chirp_v2/mert_25_2x4k.npy",
    device="cpu",
)

USE_COMPILE = False
gpt_ckpt_path = _get_model_if_needed(
    "/app/suno/checkpoints/2024-12-08_06-09-46/last_ckpt_infer.pt"
)
engine = Engine(
    gpt_ckpt_path,
    "/app/suno/models/chirp_v2/tokenizer_60k.json",
    max_sequences=MAX_STREAMS,
    compile=USE_COMPILE,
)
model = engine.model
cfg = model.config
tokenizer = engine.tokenizer

audio = load_audio("/home/sara/samples/v4_echoes_orbit.mp3").get_segment(
    from_s=0, to_s=10
)
audio.play()
in_infill_arr = process_audio(audio, cfg)

diffusion_gen.preload_models(compile=True)
engine = UpsampleEngine(min_chunk_size=25 * 5)
diffusion_generation_config = DiffusionGenerationConfig()
audio_semantic_codes = torch.from_numpy(in_infill_arr[:, 0]).reshape(1, -1).long()
print(audio_semantic_codes.shape)

for i in range(20):
    print(f"Generation {i+1}")
    engine.run_request(
        Request(
            id="0",
            generation_config=diffusion_generation_config,
            tokens=[c for c in audio_semantic_codes[0]],
            input_tokens_finished=True,
        )
    )