# %%
import os
from tqdm import tqdm
from suno_utils.utils.text import read_jsonl, write_jsonl
from suno_utils.utils.s3 import read_from_s3, list_s3_dir
import gc
import numpy as np

# %%
OUT_DATA_DIR = "/home/sara/sfx/v0"

# %%
bundle_metas = read_jsonl("/app2/suno/data/diffusion/sfx/v0/combined_metas.jsonl")

# %%
for row in bundle_metas:
    dur = row["duration_s"]
    row["tags"].append(f"duration_s: {dur}")

# %%
for k, v in bundle_metas[0].items():
    print(f"{k}: {v}")

# %%
filepath = "s3://suno-data/datasets/bundles/v5/sfx_all/dac_vae_tuned_25hz/part_0.npz"

data = read_from_s3(filepath, read_f=np.load)
print(data.keys())

# %% [markdown]
# # Memamp

# %%
# create metas map
print(len(bundle_metas))
metas_map = {meta["id"]: meta for meta in bundle_metas}
print(len(metas_map))

# %%
SEMANTIC_RATE_HZ = 25
CHUNK_SIZE_S = 15
CHUNK_SIZE = int(CHUNK_SIZE_S * SEMANTIC_RATE_HZ)

# %%
# first get a list of all the npz parts in the s3 bucket
s3_bucket = "s3://suno-data/datasets/bundles/v5/sfx_all/dac_vae_tuned_25hz"

part_filepaths = list_s3_dir(s3_bucket)
# get all the part names
part_names = [part_filepath[0].split("/")[-1].split(".")[0] for part_filepath in part_filepaths]
unique_part_names = list(set(part_names))
print(len(unique_part_names))

# %%
idx = 1
npz_filepath = f"s3://suno-data/datasets/bundles/v5/sfx_all/dac_vae_tuned_25hz/part_{idx}.npz"

vae_data = read_from_s3(npz_filepath, read_f=np.load)
meta_ids = vae_data.keys()
for meta_id in meta_ids:
    print(meta_id)
    meta_vae_data = vae_data[meta_id]
    metadata = metas_map[meta_id]
    print(metadata)
    print(meta_vae_data.shape)
    break

# %%
# load silence vae
silence_vae = np.load("/home/sara/dac_vae_tuned_25hz_15s_silence.npz")["vae_data"]
print(silence_vae.shape)

# %%
test_split = 0.025
test_idx = int(test_split * len(unique_part_names)) - 1
val_start = 0
val_end = test_idx
train_start = val_end + 1
train_end = len(unique_part_names)

val_parts = (val_start, val_end)
tr_parts = (train_start, train_end)
print(f"Validation range {val_parts}")
print(f"Train range {tr_parts}")

# %%
val_parts = (val_start, val_end)
tr_parts = (train_start, train_end)

use_local = True

for dset_type in ["val", "tr"]:
    dset_filepaths = []
    if dset_type == "val":
        part_indices = range(val_parts[0], val_parts[1])
    else:
        part_indices = range(tr_parts[0], tr_parts[1])

    metas = []

    out_mm_vae_filepath = os.path.join(OUT_DATA_DIR, f"data_vae_{dset_type}.bin")
    # out_mm_semantic_filepath = os.path.join(OUT_DATA_DIR, f"data_semantic_{dset_type}.bin")
    out_metas_filepath = os.path.join(OUT_DATA_DIR, f"metas_{dset_type}.jsonl")

    n_offs_v = 0
    n_offs_s = 0
    to_write_len_v = 0
    to_write_len_s = 0
    total_hours = 0  # Counter for total hours of audio

    out_mm_vae = np.memmap(out_mm_vae_filepath, dtype=np.float16, mode="w+", shape=(1,))

    # clear the metas file
    with open(out_metas_filepath, "w") as f:
        f.write("")

    # Create a tqdm progress bar with hours counter
    pbar = tqdm(part_indices)
    pbar.set_description("Hours: 0.00")

    for idx, part_idx in enumerate(pbar):
        if use_local:
            npz_filepath = f"/app2/suno/data/sara/sfx/combined/dac_vae_tuned_25hz/part_{idx}.npz"
            try:
                vae_data = np.load(npz_filepath)
            except Exception as e:
                print(e)
                continue
        else:
            npz_filepath = (
                f"s3://suno-data/datasets/bundles/v5/sfx_all/dac_vae_tuned_25hz/part_{idx}.npz"
            )
            try:
                vae_data = read_from_s3(npz_filepath, read_f=np.load)
            except Exception as e:
                print(e)
                continue

        # first collect the stuff we will write to disk
        # here we have to iterate over each meta_id in vae_data
        to_write_len_v = 0
        chunks = []
        for meta_id in vae_data.keys():
            meta_vae_data = vae_data[meta_id].astype(np.float16)
            metadata = metas_map[meta_id]
            vae_chunk = meta_vae_data[:CHUNK_SIZE, :]
            n_vae_tokens = vae_chunk.shape[0]

            # lets append silence vae data to the end of the chunk
            pad_len = CHUNK_SIZE - n_vae_tokens
            cropped_silence_vae_data = silence_vae[:pad_len, :]
            vae_chunk = np.concatenate([vae_chunk, cropped_silence_vae_data], axis=0)

            to_write_len_v += vae_chunk.size

            # create a new meta
            new_meta = {
                "id": meta_id,
                "duration_s": 15.0,
                "original_duration_s": meta_vae_data.shape[0] / SEMANTIC_RATE_HZ,
                "n_vae_tokens": n_vae_tokens,
                "tags": metadata["tags"],
                "dataset": metadata["dataset"],
                "s3_filepath": metadata["s3_filepath"],
            }
            metas.append(new_meta)

            chunks.append(vae_chunk.reshape(-1))

            # Add to total hours counter
            audio_duration_hours = (meta_vae_data.shape[0] / SEMANTIC_RATE_HZ) / 3600
            total_hours += audio_duration_hours

            # Update progress bar description with current total hours
            pbar.set_description(f"Hours: {total_hours:.2f}")

        # now write to disk
        out_mm_vae = np.memmap(
            out_mm_vae_filepath,
            dtype=np.float16,
            mode="r+",
            shape=(n_offs_v + to_write_len_v,),
        )

        # convert chunks to single array
        vae_chunks = np.concatenate(chunks, axis=0)

        # convert vae_chunk to float16
        out_mm_vae[n_offs_v : n_offs_v + vae_chunks.size] = vae_chunks

        n_offs_v += vae_chunks.size

    print(f"Total hours of audio added: {total_hours:.2f} for {dset_type} set")

    print(len(metas))
    write_jsonl(metas, os.path.join(out_metas_filepath), do_append=True)

    out_mm_vae.flush()
    del out_mm_vae, f
    gc.collect()