import numpy as np
import os
import tqdm
from collections import defaultdict
import pandas as pd
from suno_utils.utils.text import write_jsonl, write_json

SEMANTIC_N_CODEBOOKS = 1
COARSE_RATE_HZ = 25
SEMANTIC_CODEBOOK_SIZE = 4000
SEMANTIC_PAD_TOKEN = SEMANTIC_CODEBOOK_SIZE
SEMANTIC_RATE_HZ = 25
COARSE_CODEBOOK_SIZE = 2048
COARSE_N_CODEBOOKS = 12
COARSE_PAD_TOKEN = COARSE_CODEBOOK_SIZE
N_TOKENS_AUDIO = 6016  # max 240s of audio

def reshift(arr):
    sem_start_idx = 0
    sem_end_idx = len(arr) - 1
    semantic_arr = arr[:, :SEMANTIC_N_CODEBOOKS]
    coarse_arr = arr[:, SEMANTIC_N_CODEBOOKS:]

    coarse_start_idx = int(round(sem_start_idx * COARSE_RATE_HZ / SEMANTIC_RATE_HZ))
    coarse_end_idx = int(round(sem_end_idx * COARSE_RATE_HZ / SEMANTIC_RATE_HZ))
    assert sem_end_idx >= 0 and coarse_start_idx >= 0
    assert not (sem_end_idx > len(semantic_arr) or coarse_end_idx > len(coarse_arr))

    # get array segments
    arr_s = semantic_arr[sem_start_idx:sem_end_idx, :].copy()
    arr_c = coarse_arr[coarse_start_idx:coarse_end_idx, :].copy()
    assert arr_s.max() <= SEMANTIC_PAD_TOKEN
    assert arr_c.max() <= COARSE_PAD_TOKEN
    assert len(arr_s) == len(arr_c)
    # concat and stack
    if len(arr_c) < N_TOKENS_AUDIO:
        arr_c = np.pad(
            arr_c,
            ((0, N_TOKENS_AUDIO - len(arr_c)), (0, 0)),
            constant_values=COARSE_PAD_TOKEN,
            mode="constant",
        )
        arr_s = np.pad(
            arr_s,
            ((0, N_TOKENS_AUDIO - len(arr_s)), (0, 0)),
            constant_values=SEMANTIC_PAD_TOKEN,
            mode="constant",
        )
    arr = np.concatenate([arr_s, arr_c], axis=-1)
    arr = arr.astype(np.uint16)
    assert arr.shape == (N_TOKENS_AUDIO, SEMANTIC_N_CODEBOOKS + COARSE_N_CODEBOOKS)
    return arr


def make_dataset(input_df, output_data_dir: str, is_val=False, npz_dir="/app/suno/data/dpo/v3_npz"):
    dset_type = "val" if is_val else "tr"
    out_mmap_path = os.path.join(output_data_dir, f"data_{dset_type}.bin")
    out_metas_path = os.path.join(output_data_dir, f"meta_{dset_type}.jsonl")
    out_info_filepath = os.path.join(output_data_dir, f"info_{dset_type}.json")

    # gather the data
    _ = np.memmap(out_mmap_path, dtype=np.uint16, mode="w+", shape=(1,))
    n_offs = 0
    tot_duration_dict = defaultdict(int)
    datasets_info = defaultdict(dict)
    n = 0
    for i, row in tqdm.tqdm(input_df.iterrows(), total=len(input_df)):
        # we need to alternate between preference: neg, pos
        # print(i, row)
        assert row["preference"] == (i % 2 == 1)
        # make mmap -- two different paths
        # upsample_npz_dir = "/app/suno/data/dpo/7b_upsample_npz"
        # resampled_npz_dir = "/app/suno/data/dpo/v3_npz"
        # normal_npz_dir = "/app/suno/data/dpo/v3_npz"
        # upsample_file = f"{upsample_npz_dir}/{row['s3_id']}.npz"
        # if os.path.exists(upsample_file) and os.stat(upsample_file).st_size > 2000:
        #     local_path = upsample_file
        # else:
        #     # positive is resampled
        #     # negative is not :)
        #     if row["preference"] is True:
        #         local_path = f"{resampled_npz_dir}/{row['s3_id']}.npz"
        #     else:
        #         local_path = f"{resampled_npz_dir}/{row['s3_id']}.npz"
        local_path = (
            f"{npz_dir}/{row['s3_id']}.npz"
            if row["is_7b"]
            else f"/app/suno/data/dpo/7b_npz/{row['s3_id']}.npz"
        )
        if not os.path.exists(local_path):
            # print(row, local_path)
            raise ValueError()
        # print(local_path)
        # print( np.load(local_path))
        try:
            arr = (
                np.load(local_path)["v3.0_raw"]
                if row["is_7b"]
                else np.load(local_path)["v2_raw"]
            )
        except Exception as e:
            print(local_path)
            raise e
        assert arr.shape[0] <= 3000
        assert arr.shape[1] == 13
        arr_duration = arr.shape[0] / 25
        # print(arr.shape)
        arr = reshift(arr)
        # print("after shift and pad", arr.shape)
        arr = arr.reshape(
            -1,
        )
        # print(arr.shape)
        out_mm = np.memmap(
            out_mmap_path,
            dtype=np.uint16,
            mode="r+",
            shape=(n_offs + arr.size,),
        )
        out_mm[n_offs : n_offs + arr.size] = arr
        # print(f"offset is: {n_offs}")
        # break
        # write it once
        out_mm.flush()
        del out_mm

        add_metas = []
        add_meta = {
            "dataset": f"perference_{int(row['preference'])}",
            "id": row["s3_id"],  # this is the row s3_id
            "start_s": row["total_start_s"] if row["total_start_s"] >= 0 else None,
            "end_s": (
                row["total_clip_s"] if row["total_clip_s"] >= 0 else None
            ),  # for full clips we do know it has an edding, other wise, we don't know
            "original_duration_s": arr_duration,  # this needs to be... a bit more complicated, only works with concat!
            "vocal_start_s": None,  # these are unfortunately missing for now
            "vocal_end_s": None,  # these are unfortunately missing for now
            "tags": [
                row["tags"] if not pd.isna(row["tags"]) else ""
            ],  # tags is a list, do you know :)
            "text": row["prompt"] if not pd.isna(row["prompt"]) else "",
        }
        add_metas.append(add_meta)
        tot_duration_dict[row["preference"]] += arr_duration
        write_jsonl(
            add_metas,
            os.path.join(out_metas_path),
            do_append=bool(n_offs != 0),
        )
        if "idx_list" not in datasets_info[add_meta["dataset"]]:
            datasets_info[add_meta["dataset"]]["idx_list"] = [n]
        else:
            datasets_info[add_meta["dataset"]]["idx_list"].append(n)
        n += 1
        n_offs += arr.size

    write_json(datasets_info, out_info_filepath)
    print(f"Total {n} clips")
    for k, v in tot_duration_dict.items():
        print(f"{round(v / 60 / 60):,} hours of {k}")
    print("Done")