import os
import shutil
import torch
import numpy as np
import pandas as pd

from tqdm import tqdm
from suno_utils.utils.text import write_jsonl, read_jsonl

if __name__ == "__main__":
    # --------------------------------------------------------------------------
    # Constants
    VAE_MEMMAP_SIZE = 750
    SEMANTIC_MEMMAP_SIZE = 750
    VAE_DIM = 128
    CHUNK_SIZE = 100

    # set dataset type
    DSET_TYPE = "val"

    # load pkl file
    DATA_PKL_PATH = (
        "/home/tony/Data/Preference/up_v1/interesting_clips_up_v1_20241118_full.pkl"
    )
    NPZ_DIR = "/app/suno/data/dpo/diff_v1"
    OUT_DATA_DIR = "/app/suno/data/diffusion_ft/interesting_clips_up_v1_20241118_full"

    # --------------------------------------------------------------------------

    # create output dir
    if os.path.exists(OUT_DATA_DIR):
        # delete it
        shutil.rmtree(OUT_DATA_DIR)
    os.makedirs(OUT_DATA_DIR, exist_ok=True)

    df = pd.read_pickle(DATA_PKL_PATH)
    print(f"df shape: {df.shape}")

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

    # initial write
    out_mm_semantic = np.memmap(
        out_mm_semantic_filepath,
        dtype=np.uint16,
        mode="w+",
        shape=(1),
    )
    out_mm_vae = np.memmap(
        out_mm_vae_filepath,
        dtype=np.float16,
        mode="w+",
        shape=(1),
    )

    n_offs_s = 0
    n_offs_v = 0

    # create indices for the df
    # we will count by 2, so take all pairs of rows starting at index 0
    df_indices = list(range(0, len(df), 2))

    # now chunk the df indices
    df_indices_chunks = [
        df_indices[i : i + CHUNK_SIZE] for i in range(0, len(df_indices), CHUNK_SIZE)
    ]

    print(f"total chunks: {len(df_indices_chunks)}")

    # two consecutive rows form the positive and negative pair
    for chunk_idx, df_indices in enumerate(tqdm(df_indices_chunks)):
        arr_s_list = []
        arr_v_list = []
        new_metas = []

        for i in df_indices:
            positive_row = df.iloc[i]
            negative_row = df.iloc[i + 1]

            # load npz files
            positive_npz_path = os.path.join(NPZ_DIR, f"{positive_row['id_x']}.npz")
            negative_npz_path = os.path.join(NPZ_DIR, f"{negative_row['id_x']}.npz")

            for id_x in [positive_row["id_x"], negative_row["id_x"]]:
                # get semantic codes
                coarse_npz_path = os.path.join(NPZ_DIR, f"{id_x}.npz")
                coarse_npz = np.load(coarse_npz_path)
                codes = coarse_npz["v3.0_raw"]
                semantic_codes = codes[:, 0].astype(np.uint16)
                arr_s = semantic_codes[:SEMANTIC_MEMMAP_SIZE]

                # get vae latents
                vae_npz_path = os.path.join(NPZ_DIR, f"{id_x}_vae.npz")
                vae_npz = np.load(vae_npz_path)
                vae_latents = vae_npz["vae_latents"]
                arr_v = vae_latents[:VAE_MEMMAP_SIZE, :].astype(np.float16)

                # check for nan in arr_v or arr_s
                if not np.all(np.isfinite(arr_v)) or not np.all(np.isfinite(arr_s)):
                    continue

                # create a new meta
                meta = {
                    "text": positive_row["metadata"]["prompt"],
                    "tags": positive_row["metadata"]["tags"],
                    "n_vae_tokens": VAE_MEMMAP_SIZE,
                    "start_s": 0.0,
                    "end_s": 30.0,
                }

                try:
                    if arr_s.size < SEMANTIC_MEMMAP_SIZE:
                        continue
                    if arr_v.size < VAE_MEMMAP_SIZE * VAE_DIM:
                        continue
                except Exception as e:
                    print(f"error loading {meta['id']}: {e}")
                    continue

                arr_s_list.append(arr_s)
                arr_v_list.append(arr_v)
                new_metas.append(meta)

        print(len(arr_s_list), len(arr_v_list), len(new_metas))
        assert len(arr_s_list) == len(arr_v_list) == len(new_metas)

        to_write_len_s = SEMANTIC_MEMMAP_SIZE * len(arr_s_list)
        to_write_len_v = VAE_MEMMAP_SIZE * VAE_DIM * len(arr_v_list)

        out_mm_semantic = np.memmap(
            out_mm_semantic_filepath,
            dtype=np.uint16,
            mode="r+",
            shape=(n_offs_s + to_write_len_s,),
        )

        out_mm_vae = np.memmap(
            out_mm_vae_filepath,
            dtype=np.float16,
            mode="r+",
            shape=(n_offs_v + to_write_len_v,),
        )

        # write to memmap (has to happen sequentially)
        for new_meta, arr_s, arr_v in zip(new_metas, arr_s_list, arr_v_list):
            out_mm_semantic[n_offs_s : n_offs_s + arr_s.size] = arr_s.reshape(
                -1,
            )
            out_mm_vae[n_offs_v : n_offs_v + arr_v.size] = arr_v.reshape(
                -1,
            )
            n_offs_s += arr_s.size
            n_offs_v += arr_v.size

        # write it once
        out_mm_semantic.flush()
        out_mm_vae.flush()
        del out_mm_semantic, out_mm_vae

        write_jsonl(new_metas, out_metas_filepath, do_append=True)