import os
import gc
import re
import tqdm
import math
import copy
import json
import funcy
import random
import torch
import tempfile
import fasttext
import numpy as np
import collections
import pandas as pd
from joblib import Parallel, delayed

from collections import Counter
from transformers import PreTrainedTokenizerFast
from transformers import BertTokenizerFast
from tokenizers import (
    decoders,
    models,
    normalizers,
    pre_tokenizers,
    processors,
    trainers,
    Tokenizer,
)


from suno_utils.audio import Audio
from suno_utils.utils.tokenizers import tokenize
from suno_utils.utils.lyrics import remove_speakers
from suno_utils.utils.display import capture_output
from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists, open_from_s3
from suno_utils.harvest.youtube.constants.text_lang import BASE_TO_FASTTEXT_REMAP
from suno_utils.utils.text import (
    write_jsonl,
    read_jsonl,
    write_json,
    read_json,
    normalize_whitespace,
)

# prep metas for each dataset with one row for each entry and (eg non-music will be filtered out)
# genius_hq, youtube_music, freesound, jamendo, imslp, deezer, ytm_tagged

# metas will contain text and tags:
#   text will be a simple string (eg for foreign etc)
#   text_segments will be a list of timestamped segments for aligned
#   tags will be (maybe empty) list of strings which could be any description (genre, title etc)

# save this all to METAS_DIR

LANG_ID_MODEL_FP = "s3://suno-data/georg/trained_models/chirp_v1/lid.176.bin"
with capture_output():
    text_lang_model = read_from_s3(LANG_ID_MODEL_FP, read_f=fasttext.load_model)


def _get_text_lang(text):
    """get probability of input language for text"""
    text = text.replace("’", "'").lower()
    text = re.sub(r"\[.+?\]", " ", text)
    text = normalize_whitespace(text)
    out = text_lang_model.predict(text, k=1)
    lang_str = out[0][0]
    p_lang = out[1][0]
    lang = lang_str.split("__")[-1]
    lang = BASE_TO_FASTTEXT_REMAP.get(lang, lang)
    #     if p_lang >= 0.8:
    #         return lang
    return p_lang, lang


def _clean_tag(s):
    if s is None or not isinstance(s, str):
        return ""
    # remove symbold with special meaning
    s = re.sub(r"[\{\}\_\[\]]", " ", s)
    # squash whitespace
    return normalize_whitespace(s)


# -------- functions for parsing dataset metadata ------------


def process_imslp(base_metas_imslp: list):
    # base_metas_imslp = read_from_s3(
    #    "s3://suno-data/datasets/bundles/v1/imslp/metas.jsonl", read_f=read_jsonl
    # )

    print(f"{len(base_metas_imslp):,} clips")
    blocked_tags = ["vocal", "chorus"]
    blocked_re = re.compile(
        r"(" + r")|(".join([s.lower() for s in blocked_tags]) + r")"
    )
    allowed_ids = set()
    total_duration = 0
    durations = []
    for m in base_metas_imslp:
        tag_str = (
            m.get("title", "")
            + " "
            + m.get("track_title", "")
            + " "
            + m.get("composer", "")
            + " "
            + m.get("language", "")
            + " "
            + m.get("recording_category", "")
            + " "
            + m.get("genre", "")
            + " "
            + m.get("instruments", "")
        ).lower()
        total_duration += m["duration_s"]
        durations.append(m["duration_s"])
        if blocked_re.search(tag_str):
            continue
        if len(normalize_whitespace(tag_str)) == 0:
            continue
        allowed_ids.add(m["id"])
    print(f"{len(allowed_ids):,} music clips")
    print(f"{total_duration / 3600 / 1000:0.2f}k hr")
    # 278,620 clips
    # 274,275 music clips

    genre_counter = collections.Counter()
    instruments_counter = collections.Counter()
    composer_counter = collections.Counter()
    for m in base_metas_imslp:
        genre_counter[m.get("genre", "").lower()] += 1
        instruments_counter[m.get("instruments", "").lower()] += 1
        composer_counter[m.get("composer", "").lower()] += 1

    print(genre_counter)
    block_list = set(
        [
            "simpson, daniel léo",
            "sousa, john philip",
            "various",
            "zhang, shuwen",
            "bartók, béla",
            "strauss, richard",
            "prokofiev, sergey",
        ]
    )
    for composer in composer_counter:
        if composer_counter[composer] > 500:
            print(composer)

    composers_and_death = """
    monteverdi, claudio - 1643
    purcell, henry - 1695
    corelli, arcangelo - 1713
    bach, johann sebastian - 1750
    telemann, georg philipp - 1767
    scarlatti, domenico - 1757
    handel, george frideric - 1759
    mozart, wolfgang amadeus - 1791
    beethoven, ludwig van - 1827
    bach, carl philipp emanuel - 1788
    gluck, christoph willibald - 1787
    haydn, joseph - 1809
    sor, fernando - 1839
    donizetti, gaetano - 1848
    schumann, robert - 1856
    rossini, gioacchino - 1868
    liszt, franz - 1886
    wagner, richard - 1883
    brahms, johannes - 1897
    verdi, giuseppe - 1901
    sullivan, arthur - 1900
    debussy, claude - 1918
    granados, enrique - 1916
    mahler, gustav - 1911
    puccini, giacomo - 1924
    satie, erik - 1925
    franck, césar - 1890
    tchaikovsky, pyotr - 1893
    gounod, charles - 1893
    prokofiev, sergey - 1953
    rachmaninoff, sergei - 1943
    scriabin, aleksandr - 1915
    paganini, niccolò - 1840
    elgar, edward - 1934
    villa-lobos, heitor - 1959
    berlioz, hector - 1869
    massenet, jules - 1912
    rimsky-korsakov, nikolay - 1908
    offenbach, jacques - 1880
    franz lehár - 1948
    saint-saëns, camille - 1921
    sibelius, jean - 1957
    franz schubert - 1828
    manuel de falla - 1946
    george gershwin - 1937
    felix mendelssohn - 1847
    michel rondeau - Unknown
    johann strauss jr. - 1899
    john philip sousa - 1932
    vincenzo bellini - 1835
    georges bizet - 1875
    georg philipp telemann - 1767
    antonín dvořák - 1904
    modest mussorgsky - 1881
    antonio vivaldi - 1741
    camille saint-saëns - 1921
    jean sibelius - 1957
    pierre-montan berton - 1780
    bryan d. hoyt - Unknown
    claude le jeune - 1600
    alexander nakarada - Unknown
    """
    composers = []
    allowed_composers = set()
    for l in composers_and_death.split("\n"):
        outs = l.split("-")
        if len(outs) < 2:
            continue
        composer = outs[0].strip()
        date = outs[1].strip()
        if not date.isnumeric():
            continue
        if int(date) <= 1923:
            composers.append((int(date), composer))
            allowed_composers.add(composer)
    print(sorted(allowed_composers))

    # build metadata file
    metas = []
    seen_ids = set()
    total_duration = 0
    for m in base_metas_imslp:
        if m["id"] in seen_ids or m["id"] not in allowed_ids:
            continue
        if m.get("composer", "").lower() not in allowed_composers:
            continue
        new_m = {"id": m["id"]}
        # original tags
        spaced_tags = []
        for t in [
            m.get("genre", ""),
            m.get("title", ""),
            m.get("instruments", ""),
        ]:
            spaced_tags.extend(t.split(", "))
        composer_tag = m.get("composer", "")
        if "," in composer_tag:
            composer_tag = " ".join(composer_tag.split(", ")[::-1])
        spaced_tags.append(composer_tag)
        tags = []
        for tag in spaced_tags:
            tag = tag.lower()
            no_pattern = r"no\. (\d+)"
            op_pattern = r"op\. (\d+)"
            bwv_pattern = r"bwv (\d+)"
            k_pattern = r"k\. (\d+)"
            l_pattern = r"l\. (\d+)"
            tag = re.sub(no_pattern, "", tag)
            tag = re.sub(op_pattern, "", tag)
            tag = re.sub(bwv_pattern, "", tag)
            tag = re.sub(k_pattern, "", tag)
            tag = re.sub(l_pattern, "", tag)
            if "voice" in tag.lower():
                continue
            if "chorus" in tag.lower():
                continue
            if len(tag.strip()) == 0:
                continue
            tag = tag.replace("  ", " ")
            tags.append(tag)
        if len(tags) > 0:
            new_m["tags"] = list(set(tags))
        private_tags = [
            _clean_tag(m.get("title")),
            _clean_tag(m.get("track_title", "")),
            _clean_tag(m.get("composer", "")),
        ]
        spaced_tags = []
        for t in private_tags:
            spaced_tags.extend(t.split(", "))
        private_tags = spaced_tags
        private_tags += tags
        private_tags = [t for t in private_tags if len(t) > 0]
        if len(private_tags) > 0:
            new_m["private_tags"] = list(set(private_tags))

        new_m["s3_filepath"] = m.get("s3_filepath")
        metas.append(new_m)
        seen_ids.add(m["id"])
        total_duration += m["duration_s"]

    print(f"{len(metas):,} entries, total duration {total_duration/60/60:,.0f} hr")
    return metas


def process_jamendo(base_metas: str):
    # base_metas = read_from_s3(
    #    "s3://suno-data/datasets/bundles/v1/jamendo/metas.jsonl", read_f=read_jsonl
    # )

    metas = []
    seen_ids = set()
    for m in base_metas:
        if m["id"] in seen_ids:
            continue
        new_m = {"id": m["id"], "s3_filepath": m["audio_filepath"]}

        tags = []
        tags.extend([_clean_tag(s) for s in m.get("genre", "").split(";")])
        tags.extend([_clean_tag(s) for s in m.get("instrument", "").split(";")])
        tags.extend([_clean_tag(s) for s in m.get("mood/theme", "").split(";")])
        tags = [t for t in tags if len(t) > 0]
        if len(tags) > 0:
            new_m["tags"] = list(set(tags))

        metas.append(new_m)
        seen_ids.add(m["id"])

    print(f"{len(metas):,} entries")
    # 55,609 entries
    return metas


def process_genius_hq(
    base_metas: str, alignment_metas: str = None, audio_production_metas: str = None
):

    if alignment_metas is not None:
        aligned_lyrics_map = {k: v for k, v in alignment_metas}
        del alignment_metas
        gc.collect()
    else:
        aligned_lyrics_map = None

    if audio_production_metas is not None:
        audio_production_map = {k: v for k, v in audio_production_metas}
        del audio_production_metas
        gc.collect()
    else:
        audio_production_map = None

    metas = []
    seen_ids = set()

    for m in tqdm.tqdm(base_metas):
        if (
            m["genius_views"] < 50
            or m["youtube_views"] < 100  # default is 50
            or m["duration_s"] < 1 * 60  # default is 100
            or m["duration_s"] > 8 * 60
            or len(m["lyrics"]) > 6144
            or len(m["lyrics"]) < 50
        ):
            continue

        if audio_production_map is not None:
            if m["id"] in audio_production_map:
                continue

        if m["id"] in seen_ids:
            continue

        new_m = {
            "id": m["id"],
            "original_id": m["genius_slug"],
            "views": m["youtube_views"],
            "lang": m["lang"].split("-")[0].lower(),
            "dataset": "genius_hq",
        }
        text_segments = []
        for mm in aligned_lyrics_map[m["id"]]:
            text_segments.append(
                {
                    "text": remove_speakers(mm["text"]),
                    "private_text": mm["text"],
                    "start_s": mm["start_s"],
                    "end_s": mm["end_s"],
                    "vocal_start_s": mm["vocal_start_s"],
                    "vocal_end_s": mm["vocal_end_s"],
                }
            )
        new_m["text_segments"] = text_segments
        new_m["duration_s"] = m["duration_s"]
        new_m["private_text"] = m["lyrics"]
        new_m["text"] = remove_speakers(m["lyrics"])
        new_m["s3_filepath"] = m["audio_filepath"]
        # add tags
        tags = [
            normalize_whitespace(t_clean.replace("Genius", " "))
            for t in m.get("tags_text", [])
            if len(t_clean := _clean_tag(t)) > 0
        ]
        if len(tags) == 0:
            # we don't want to fine tune on things with no tags
            continue
        if len(tags) > 0:
            new_m["tags"] = list(set(tags))

        private_tag = _clean_tag(m["youtube_title"])
        private_tags = tags
        if len(private_tag) > 0:
            private_tags += [private_tag]
        if len(private_tags) > 0:
            new_m["private_tags"] = list(set(private_tags))
        metas.append(new_m)
        seen_ids.add(m["id"])

    print(f"{len(metas):,} entries")
    return metas


def process_youtube_music(base_metas: list, alignment_metas: str = None):

    print(f"{len(base_metas):,} clips")
    allowed_ids = set()
    for m in base_metas:
        if m["view_count"] < 1000:
            continue
        if re.search(r"\blive\b", m["title"], flags=re.IGNORECASE):
            continue
        allowed_ids.add(m["id"])
    print(f"{len(allowed_ids):,} music clips")
    # 2,101,712 clips
    # 2,013,545 music clips

    # load aligned lyrics
    if alignment_metas is not None:
        aligned_lyrics_map = {k: v for k, v in alignment_metas}
        del alignment_metas
        gc.collect()

        metas = []
        seen_ids = set()
        for m in base_metas:
            if (
                m["id"] in seen_ids
                or m["id"] not in allowed_ids
                or m["id"] not in aligned_lyrics_map
            ):
                continue
            new_m = {
                "id": m["id"],
                "views": m["view_count"],
                "s3_filepath": m["s3_filepath"],
            }
            # get tags
            tags = [
                m.get("genre_guess"),
                m.get("mood_guess"),
            ]
            tags = [t_clean for t in tags if len(t_clean := _clean_tag(t)) > 0]
            if len(tags) > 0:
                new_m["tags"] = list(set(tags))
            # get private tags
            private_tags = [
                m.get("title"),
                m.get("artists"),
                m.get("album"),
            ]
            private_tags = [
                t_clean for t in private_tags if len(t_clean := _clean_tag(t)) > 0
            ]
            private_tags += tags
            if len(private_tags) > 0:
                new_m["private_tags"] = list(set(private_tags))
            # get aligned lyrics
            text_segments = []
            for mm in aligned_lyrics_map[m["id"]]:
                text_segments.append(
                    {
                        "text": remove_speakers(mm["text"]),
                        "private_text": mm["text"],
                        "start_s": mm["start_s"],
                        "end_s": mm["end_s"],
                        "vocal_start_s": mm["vocal_start_s"],
                        "vocal_end_s": mm["vocal_end_s"],
                    }
                )
            new_m["text_segments"] = text_segments

            # get lyrics
            if (
                "lang_guess" in m
                and "lyrics" in m
                and
                #         m["lang_guess"].lower() != "en" and
                #         _get_text_lang(m["lyrics"])[1] != "en" and
                #         len(m["lang_guess"].lower().split("-")[0]) >= 2 and
                len(m["lyrics"]) <= 5120  # empirically verified
            ):
                use_lyrics = False
                if m["lang_guess"][:2] in ("zh", "ja"):
                    if len(m["lyrics"]) < 256:
                        use_lyrics = True
                else:
                    if len(m["lyrics"]) < 512:
                        use_lyrics = True
                if use_lyrics:
                    lang_guess_1 = m["lang_guess"].split("-")[0].lower()
                    lang_guess_2 = _get_text_lang(m["lyrics"])[1].split("-")[0].lower()
                    if (lang_guess_1 == "en") == (lang_guess_2 == "en"):
                        # make sure they agree on english or not
                        lyrics = m["lyrics"]
                        new_m["text"] = m["lyrics"]
                        new_m["lang"] = lang_guess_1
            metas.append(new_m)
            seen_ids.add(m["id"])
        print(f"{len(metas):,} entries")
        print(f"{len([m for m in metas if 'text' in m]):,} entries with lyrics")
        # 2,013,545 entries
        # 674,466 entries with lyrics
    else:
        metas = []
        seen_ids = set()
        for m in base_metas:
            if m["id"] in seen_ids or m["id"] not in allowed_ids:
                continue
            new_m = {
                "id": m["id"],
                "views": m["view_count"],
                "s3_filepath": m["s3_filepath"],
            }
            # get tags
            tags = [
                m.get("genre_guess"),
                m.get("mood_guess"),
            ]
            tags = [t_clean for t in tags if len(t_clean := _clean_tag(t)) > 0]
            if len(tags) > 0:
                new_m["tags"] = list(set(tags))
            # get private tags
            private_tags = [
                m.get("title"),
                m.get("artists"),
                m.get("album"),
            ]
            private_tags = [
                t_clean for t in private_tags if len(t_clean := _clean_tag(t)) > 0
            ]
            private_tags += tags
            if len(private_tags) > 0:
                new_m["private_tags"] = list(set(private_tags))

            # get lyrics
            if (
                "lang_guess" in m
                and "lyrics" in m
                and
                #         m["lang_guess"].lower() != "en" and
                #         _get_text_lang(m["lyrics"])[1] != "en" and
                #         len(m["lang_guess"].lower().split("-")[0]) >= 2 and
                len(m["lyrics"]) <= 5120  # empirically verified
            ):
                use_lyrics = False
                if m["lang_guess"][:2] in ("zh", "ja"):
                    if len(m["lyrics"]) < 256:
                        use_lyrics = True
                else:
                    if len(m["lyrics"]) <= 5120:
                        use_lyrics = True
                # if len(m["lyrics"]) < 512:
                #    use_lyrics = True
                # else:
                #    print(len(m["lyrics"]))
                if use_lyrics:
                    lang_guess_1 = m["lang_guess"].split("-")[0].lower()
                    lang_guess_2 = _get_text_lang(m["lyrics"])[1].split("-")[0].lower()
                    if (lang_guess_1 == "en") == (lang_guess_2 == "en"):
                        # make sure they agree on english or not
                        lyrics = m["lyrics"]
                        new_m["text"] = m["lyrics"]
                        new_m["lang"] = lang_guess_1
            metas.append(new_m)
            seen_ids.add(m["id"])
        print(f"{len(metas):,} entries")
        print(f"{len([m for m in metas if 'text' in m]):,} entries with lyrics")
        # 2,013,545 entries
        # 674,466 entries with lyrics

    return metas


def process_freesound(metas_dir: str):
    base_metas = read_from_s3(
        "s3://suno-data/datasets/bundles/v1/freesound/metas.jsonl", read_f=read_jsonl
    )
    print(f"{len(base_metas):,} clips")
    print(f"{np.sum([m['duration_s'] for m in base_metas])/60/60:,.0f} hours of audio")
    needle_tags = [
        "music",
        "synth",
        "bpm",
        "drums",
        "drum-loop",
        "guitar",
        "piano",
        "soundtrack",
        "song",
        "violin",
        "melody",
        "bassline",
    ]

    needle_re = re.compile(
        r"(\b" + r"\b)|(\b".join([s.lower() for s in needle_tags]) + r"\b)"
    )
    l = []
    music_ids = set()
    for m in base_metas:
        if m["duration_s"] < 10:
            continue
        tag_str = " ".join(m["tags"]).lower()
        if needle_re.search(tag_str):
            l.append(m["duration_s"])
            music_ids.add(m["id"])
    print(f"{len(music_ids):,} music clips")
    print(f"{np.sum(l)/60/60:,.0f} hours of music")
    # 510,491 clips
    # 2,993 hours of audio
    # 34,225 music clips
    # 411 hours of music

    metas = []
    seen_ids = set()
    for m in base_metas:
        if m["id"] in seen_ids or m["id"] not in music_ids:
            continue
        new_m = {"id": m["id"], "s3_filepath": m["s3_filepath"]}
        tags = [
            _clean_tag(m.get("cap1")),
            _clean_tag(m.get("cap2")),
        ]
        tags.extend([_clean_tag(s) for s in m.get("tags", [])])
        tags = [t for t in tags if len(t) > 0]
        if len(tags) > 0:
            new_m["tags"] = list(set(tags))
        metas.append(new_m)
        seen_ids.add(m["id"])
    print(f"{len(metas):,} entries")
    # 34,225 entries
    return metas


def process_pond5_music(pd5_base_metas: str):
    # pd5_base_metas = read_from_s3(
    #    "s3://suno-data/datasets/bundles/v2/pond5_music/metas.jsonl", read_f=read_jsonl
    # )
    pd5_extra_metas = read_from_s3(
        "s3://suno-data/datasets/harvest/pond5_music/pond5_metas.jsonl",
        read_f=read_jsonl,
    )

    print(len(pd5_extra_metas))
    extra_info = {}
    seen_ids = set()
    genre_counter = collections.Counter()
    for m in tqdm.tqdm(pd5_extra_metas):
        # if "s3_filepath" not in m:
        #    continue
        if m["id"] in seen_ids:
            continue

        # if m["duration"] < 30:
        #    continue
        # if m["duration"] > (60 * 6):
        #    continue

        name = m.get("name", "")
        description = m.get("description", "")
        tags = m.get("tags", [])
        genre = m.get("genre", "")

        # if not genre:
        #    continue

        # genre_counter[genre.lower().strip()] += 1

        # if not description or not tags or not name:
        #    continue
        # if len(description.strip().split()) < 10:
        #    print("short description", m)
        #    continue
        # if len(description.strip().split()) > 100:
        #    continue
        # if not isinstance(tags, list) or len(tags) < 5 or len(tags) > 50:
        #    continue
        # if there are weird tags
        # if max(len(tag) for tag in tags) > 30:
        #    continue

        extra_info[m["id"]] = {
            "name": name if isinstance(name, str) else "",
            "description": (description if isinstance(description, str) else ""),
            "tags": (
                ([genre] if genre else []) + tags if isinstance(tags, list) else ""
            ),
        }
        seen_ids.add(m["id"])
    print(len(seen_ids))

    metas = []
    total_duration = 0
    n_c = 0
    for m in tqdm.tqdm(pd5_base_metas):
        #     if n_c > 5:
        #         break
        if m["id"] not in seen_ids:
            print("not in seen_ids")
            continue
        # if random.random() < 0.93:
        #    continue
        n_c += 1
        info = extra_info[m["id"]]
        new_m = {"id": m["id"], "s3_filepath": m["s3_filepath"]}
        # print(m, info)
        tags = [
            info["name"],
            info["description"],
        ]
        tags.extend(info["tags"])
        tags = [_clean_tag(t) for t in tags]
        tags = [t for t in tags if len(t) > 0]
        new_m["tags"] = list(set(tags))
        metas.append(new_m)
        total_duration += m["duration_s"]
    print(f"{len(metas):,} entries, total duration {total_duration/60/60:,.0f} hr")
    return metas


def process_deezer(base_metas: list):
    # base_metas = read_from_s3(
    #    "s3://suno-data/datasets/bundles/v2/deezer/metas_plus.jsonl", read_f=read_jsonl
    # )
    print(f"{len(base_metas):,} clips")

    prod_tag_version = "v6"

    aligned_lyrics = read_from_s3(
        f"s3://suno-data/datasets/metadata/alignments/deezer_hq_alignments_{prod_tag_version}.jsonl",
        read_f=read_jsonl,
    )
    aligned_lyrics_map = {k: v for k, v in aligned_lyrics}
    del aligned_lyrics
    gc.collect()

    allowed_ids = set()
    for m in base_metas:
        # TODO: the threshold here is low...
        if m["views"] < 50:
            continue
        #     if re.search(r"\blive\b", m["title"], flags=re.IGNORECASE):
        #         continue
        allowed_ids.add(m["id"])
    print(f"{len(allowed_ids):,} music clips")
    # 768,515 clips
    # 535,018 music clips

    metas = []
    seen_ids = set()
    for m in base_metas:
        if (
            m["id"] in seen_ids
            or m["id"] not in allowed_ids
            or m["id"] not in aligned_lyrics_map
        ):
            continue
        if (
            "lang_guess" not in m
            or len(m["lang_guess"]) < 2
            or "lyrics" not in m
            or len(m["lyrics"]) > 5120  # empirically verified
        ):
            continue
        if m["lang_guess"][:2] in ("zh", "ja"):
            if len(m["lyrics"]) < 256:
                continue
        else:
            if len(m["lyrics"]) < 512:
                continue
        new_m = {
            "id": m["id"],
            "views": m["views"],
            "text": m["lyrics"],
            "lang": m["lang_guess"].split("-")[0].lower(),
            "s3_filepath": m["s3_filepath"],
        }
        # get private tags
        private_tags = [
            m.get("youtube_title"),
            m.get("youtube_artist"),
        ]
        private_tags = [
            t_clean for t in private_tags if len(t_clean := _clean_tag(t)) > 0
        ]
        if len(private_tags) > 0:
            new_m["private_tags"] = list(set(private_tags))
        metas.append(new_m)
        seen_ids.add(m["id"])
    print(f"{len(metas):,} entries")
    # 508,804 entries
    return metas


def process_ytm_tagged(base_metas: list, alignment_metas: str = None):

    # base_metas = read_from_s3(
    #    "s3://suno-data/datasets/bundles/v2/ytm_tagged/metas.jsonl", read_f=read_jsonl
    # )

    print(f"{len(base_metas):,} clips")
    allowed_ids = set()
    for m in base_metas:
        # if m["views"] < 100:
        #    continue
        #     if re.search(r"\blive\b", m["title"], flags=re.IGNORECASE):
        #         continue
        allowed_ids.add(m["id"])
    print(f"{len(allowed_ids):,} music clips")
    # 2,772,393 clips
    # 2,255,305 music clips

    if alignment_metas is not None:
        aligned_lyrics_map = {k: v for k, v in alignment_metas}

        metas = []
        seen_ids = set()
        for m in base_metas:
            if (
                m["id"] in seen_ids
                or m["id"] not in allowed_ids
                or m["id"] not in aligned_lyrics_map
            ):
                continue
            new_m = {
                "id": m["id"],
                "views": m["views"],
                "s3_filepath": m["s3_filepath"],
            }

            text_segments = []

            for mm in aligned_lyrics_map[m["id"]]:
                text_segments.append(
                    {
                        "text": remove_speakers(mm["text"]),
                        "private_text": mm["text"],
                        "start_s": mm["start_s"],
                        "end_s": mm["end_s"],
                        "vocal_start_s": mm["vocal_start_s"],
                        "vocal_end_s": mm["vocal_end_s"],
                    }
                )

            new_m["text_segments"] = text_segments
            new_m["text"] = remove_speakers(m["lyrics"])

            # get tags
            tags = m.get("tags", [])
            tags = [t_clean for t in tags if len(t_clean := _clean_tag(t)) > 0]
            if len(tags) > 0:
                new_m["tags"] = list(set(tags))
            # get private tags
            private_tags = m.get("private_tags", [])
            private_tags = [
                t_clean for t in private_tags if len(t_clean := _clean_tag(t)) > 0
            ]
            if len(private_tags) > 0:
                new_m["private_tags"] = list(set(private_tags))
            metas.append(new_m)
            seen_ids.add(m["id"])
        print(f"{len(metas):,} entries")
        # 2,255,305 entries
    else:
        metas = []
        seen_ids = set()
        for m in base_metas:
            if m["id"] in seen_ids or m["id"] not in allowed_ids:
                continue
            new_m = {
                "id": m["id"],
                "views": m["views"],
                "s3_filepath": m["s3_filepath"],
            }

            text_segments = []

            if "lyrics" in m:
                text_segments.append(
                    {
                        "text": remove_speakers(m["lyrics"]),
                        "private_text": m["lyrics"],
                        "start_s": m["start_s"],
                        "end_s": m["end_s"],
                    }
                )
                new_m["text"] = remove_speakers(m["lyrics"])

            new_m["text_segments"] = text_segments

            # get tags
            tags = m.get("tags", [])
            tags = [t_clean for t in tags if len(t_clean := _clean_tag(t)) > 0]
            if len(tags) > 0:
                new_m["tags"] = list(set(tags))
            # get private tags
            private_tags = m.get("private_tags", [])
            private_tags = [
                t_clean for t in private_tags if len(t_clean := _clean_tag(t)) > 0
            ]
            if len(private_tags) > 0:
                new_m["private_tags"] = list(set(private_tags))
            metas.append(new_m)
            seen_ids.add(m["id"])
        print(f"{len(metas):,} entries")

    return metas


# ------------ tokenizer functions ------------


def _augment_case(s):
    if random.random() >= 0.95:
        s = s.upper()
    elif random.random() >= 0.95:
        s = s.lower()
    return s


def _augment_tag(s):
    if random.random() >= 0.95:
        s = s.upper()
    elif random.random() >= 0.95:
        s = s.capitalize()
    elif random.random() >= 0.9:
        s = s.title()
    elif random.random() >= 0.9:
        s = s.lower()
    if random.random() >= 0.5:
        s = s.replace("-", " ").strip()
    return s


def _space_repl(m):
    s = m.group()
    n_newline = s.count("\n")
    if n_newline >= 2:
        return "\n\n"
    elif n_newline == 1:
        return "\n"
    return " "


def _simplify_whitespace(text, retain_newlines=True):
    """simplify while respecting up to 2 newlines"""
    if retain_newlines:
        text = re.sub(r"\s+", _space_repl, text).strip()
    else:
        text = re.sub(r"\s+", " ", text).strip()
    return text


def make_tokenizer(base_dir: str, metas_dir: str, tmp_dir: str):

    tokenizer_train_filepath = os.path.join(tmp_dir, "tokenizer_train.txt")
    # create text file with all text
    with open(tokenizer_train_filepath, "w") as f:
        for fn in tqdm.tqdm(
            [
                "genius_hq_v6",
                "youtube_music",
                "freesound",
                "jamendo",
                "imslp",
                "pond5_music",
                "deezer",
                "ytm_tagged",
            ]
        ):
            metas_filepath = os.path.join(metas_dir, f"{fn}.jsonl")

            if not os.path.isfile(metas_filepath):
                print(f"{metas_filepath} not found. Skipping...")
                continue

            for m in read_jsonl(metas_filepath):
                text_key = (
                    "text"
                    if "private_text" not in m or random.random() >= 0.2
                    else "private_text"
                )
                if text_key in m:
                    f.write(_simplify_whitespace(_augment_case(m[text_key])) + "\n")
                text_key = (
                    "text_segments"
                    if "private_text_segments" not in m or random.random() >= 0.2
                    else "private_text_segments"
                )
                if text_key in m:
                    for mm in m[text_key]:
                        f.write(_simplify_whitespace(_augment_case(mm["text"])) + "\n")
                text_key = (
                    "tags"
                    if "private_tags" not in m or random.random() >= 0.2
                    else "private_tags"
                )
                if text_key in m:
                    join_char = random.choice([", ", "; ", " "])
                    f.write(
                        _simplify_whitespace(
                            join_char.join([_augment_tag(t) for t in m[text_key]])
                        )
                        + "\n"
                    )

    min_char_n = 3
    with open(tokenizer_train_filepath) as f:
        corpus = f.read()
    char_counts = Counter(corpus)
    print(len(char_counts), "unique chars")
    trans_table = str.maketrans(
        {k: " " for k, v in char_counts.items() if v < min_char_n}
    )
    corpus = corpus.translate(trans_table)
    char_counts_2 = Counter(corpus)
    print(len(char_counts_2), "unique chars after filtering")
    vocab = [c for c, _ in char_counts_2.most_common()]
    corpus_lines = corpus.split("\n")
    corpus_lines = [
        line_clean for line in corpus_lines if len(line_clean := line.strip()) > 0
    ]
    random.seed(6006)
    sample_corpus_lines = random.choices(corpus_lines, k=1_000_000)

    tokenizer_train_clean_filepath = os.path.join(tmp_dir, "tokenizer_train_clean.txt")

    with open(tokenizer_train_clean_filepath, "w") as f:
        for line in corpus_lines:
            f.write(line + "\n")

    # train and save tokenizer
    tokenizer = Tokenizer(models.WordPiece(unk_token="[UNK]"))
    tokenizer.normalizer = normalizers.BertNormalizer(lowercase=False)
    tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
    special_tokens = ["[UNK]", "[PAD]"]
    trainer = trainers.WordPieceTrainer(
        vocab_size=60_000, special_tokens=special_tokens
    )

    tokenizer.train([tokenizer_train_clean_filepath], trainer=trainer)
    tokenizer_path = os.path.join(base_dir, "tokenizer_60k.json")
    tokenizer.save(tokenizer_path)


if __name__ == "__main__":
    # load all metas from bundle
    base_metas = read_jsonl("metadata/diffusion_mix_v2/metas.jsonl")

    datasets = [
        "jamendo",
        "pond5_music",
        "imslp",
        "youtube_music",
        "genius_hq",
        "ytm_tagged",
    ]

    duration_mins = 0.5

    metas_dir = f"/app/suno/data/diffusion_mix_v2/metadata"
    tmp_dir = os.path.join(os.getcwd(), "tmp")
    os.makedirs(metas_dir, exist_ok=True)
    os.makedirs(tmp_dir, exist_ok=True)

    if duration_min == 2:
        use_alignemnts = False
    else:
        use_alignemnts = False

    filtered_metas = []
    for dataset in datasets:
        print(f"Processing {dataset}...")

        # filter the base metas for the specific dataset
        dataset_metas = [m for m in base_metas if m["dataset"] == dataset]
        print(f"{len(dataset_metas):,} clips in {dataset}")
        alignment_metas = None

        if dataset == "youtube_music":
            metas = process_youtube_music(
                dataset_metas, alignment_metas=alignment_metas
            )
        elif dataset == "genius_hq":
            alignment_metas = read_from_s3(
                "s3://suno-data/datasets/metadata/alignments/genius_alignments_v11.jsonl",
                read_f=read_jsonl,
            )
            audio_production_metas = read_from_s3(
                f"s3://suno-data/datasets/metadata/audio_production/genius_audio_production_v1.jsonl",
                read_f=read_jsonl,
            )
            metas = process_genius_hq(
                dataset_metas,
                alignment_metas=alignment_metas,
                audio_production_metas=audio_production_metas,
            )
        elif dataset == "ytm_tagged":
            alignment_metas = read_from_s3(
                "s3://suno-data/datasets/metadata/alignments/ytm_alignments_v11.jsonl",
                read_f=read_jsonl,
            )
            metas = process_ytm_tagged(dataset_metas, alignment_metas=alignment_metas)
        elif dataset == "pond5_music":
            metas = process_pond5_music(dataset_metas)
        elif dataset == "imslp":
            metas = process_imslp(dataset_metas)
        elif dataset == "jamendo":
            metas = process_jamendo(dataset_metas)
        else:
            raise ValueError(f"Dataset {dataset} not recognized")

        filtered_metas.extend(metas)
    write_jsonl(filtered_metas, os.path.join(metas_dir, "metas.jsonl"))