import os
import torch
import torchaudio
import numpy as np
from tqdm import tqdm

from suno_utils.utils.text import read_jsonl
from suno_utils.audio import Audio
from suno_utils.tasks.dac_2c_12cb import (
    encode as codec_encode,
    decode as codec_decode,
    EMBEDDING_RATE as CODEC_EMBEDDING_RATE,
)
from suno_utils.tasks.dac_2c_12cb import preload_models as preload_codec_models


N_TOKENS_MEMMAP = 6016
SEMANTIC_N_CODEBOOKS = 1
COARSE_N_CODEBOOKS = 12


def remove_pad_tokens(arr, value: int = 2048):
    # Find the last index where the value is not equal to the specified value
    non_value_index = np.where(arr != value)[0]

    if non_value_index.size == 0:
        # If there are no values that are not equal to the specified value, return an empty array
        return np.array([], dtype=arr.dtype)

    # Get the last index where the value is not equal to the specified value
    last_non_value_index = non_value_index[-1]

    # Slice the array to remove trailing values
    return arr[: last_non_value_index + 1]


if __name__ == "__main__":
    # configuration
    use_val = True
    tmp_dir = os.path.join(os.getcwd(), "tmp-validate")
    os.makedirs(tmp_dir, exist_ok=True)

    if use_val:
        s3_metas_filepath = "/app/suno/data/chirp_v4_ft/multi/metas_val.jsonl"
        memmap_filepath = "/app/suno/data/chirp_v4/multi/data_val.bin"
    else:
        s3_metas_filepath = "/app/suno/data/chirp_v4_ft/multi/metas_tr.jsonl"
        memmap_filepath = "/app/suno/data/chirp_v4/multi/data_tr.bin"

    # load DAC codec
    _ = preload_codec_models(
        checkpoint_filepath="s3://suno-data/georg/models/codec/dac_2c_25x12.pt",
    )

    # load memmap
    mm = np.memmap(os.path.join(memmap_filepath), dtype=np.uint16, mode="r")
    mm = mm.reshape(-1, N_TOKENS_MEMMAP, SEMANTIC_N_CODEBOOKS + COARSE_N_CODEBOOKS)

    # load the memmap aligned jsonl file
    print(f"Reading {s3_metas_filepath} from disk...")
    s3_metas = read_jsonl(s3_metas_filepath)

    print(mm.shape, len(s3_metas))
    assert mm.shape[0] == len(s3_metas)

    for idx, meta in enumerate(tqdm(s3_metas)):
        dataset = meta.get("dataset")
        s3_filepath = meta.get("s3_filepath")
        id_key = meta.get("id")

        # print(dataset, s3_filepath)
        if s3_filepath is not None:
            s3_audio = Audio.from_s3(s3_filepath, n_channels=2)
            s3_audio_sample_rate = s3_audio.sample_rate
            s3_audio = torch.from_numpy(s3_audio.array_float)
            s3_audio /= s3_audio.abs().max()

            codes = remove_pad_tokens(mm[idx, :, 1:])
            decoded_audio = codec_decode(codes)
            decoded_audio = torch.from_numpy(decoded_audio.array_float)
            decoded_audio /= decoded_audio.abs().max()

            print(s3_audio.shape, decoded_audio.shape)

            s3_audio_out = os.path.join(tmp_dir, f"{idx}-{id_key}-s3.wav")
            decoded_audio_out = os.path.join(tmp_dir, f"{idx}-{id_key}-memmap.wav")
            torchaudio.save(s3_audio_out, s3_audio, s3_audio_sample_rate)
            torchaudio.save(decoded_audio_out, decoded_audio, 48_000)

        if idx > 10:
            break