{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "import polars as pl\n",
    "from tqdm import tqdm\n",
    "\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"gpt-4_1-tagging\"\n",
    "# /home/christian/code/christian/metadata/organized/gpt-4_1-tagging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load gpt tags for discogs_subset, genius, and imslp\n",
    "discogs_path = \"/home/christian/code/christian/metadata/tagging/gpt-4_1-tagging_results_discogs_subset.json\"\n",
    "genius_path = \"/home/christian/code/christian/metadata/tagging/gpt-4_1-tagging_results_genius.json\"\n",
    "imslp_path = \"/home/christian/code/christian/metadata/tagging/gpt-4_1-tagging_results_imslp.json\"\n",
    "\n",
    "# Load all tag files and combine into a single dictionary\n",
    "gpt_tags = {}\n",
    "for path, name in [(discogs_path, \"discogs\"), (genius_path, \"genius\"), (imslp_path, \"imslp\")]:\n",
    "    with open(path) as f:\n",
    "        dataset_tags = json.load(f)\n",
    "        # for each tag, process them into a list of tags split on ';'\n",
    "        for clip_id, tags in dataset_tags.items():\n",
    "            tags = tags[\"gpt_tag\"].split(';')\n",
    "            tags = [tag.strip() for tag in tags]\n",
    "            # remove tags that are \"None\" string\n",
    "            tags = [tag for tag in tags if tag != \"None\"]\n",
    "            gpt_tags[clip_id] = tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# also load some popular tags\n",
    "discogs_subset_chart_path = \"/home/christian/code/christian/metadata/popularity/discogs_subset_chart_metas.jsonl\"\n",
    "discogs_subset_chart_metas = pl.read_ndjson(discogs_subset_chart_path)\n",
    "\n",
    "discogs_subset_grammy_path = \"/home/christian/code/christian/metadata/popularity/discogs_subset_grammy_metas.jsonl\"\n",
    "discogs_subset_grammy_metas = pl.read_ndjson(discogs_subset_grammy_path)\n",
    "\n",
    "popularity_tags = {}\n",
    "for meta in discogs_subset_chart_metas.iter_rows(named=True):\n",
    "    meta_id = meta[\"id\"]\n",
    "    # add the tags to the popularity_tags dictionary\n",
    "    if meta_id not in popularity_tags:\n",
    "        popularity_tags[meta_id] = []\n",
    "    popularity_tags[meta_id].extend(meta[\"tags\"])\n",
    "\n",
    "for meta in discogs_subset_grammy_metas.iter_rows(named=True):\n",
    "    meta_id = meta[\"id\"]\n",
    "    # add the tags to the popularity_tags dictionary\n",
    "    if meta_id not in popularity_tags:\n",
    "        popularity_tags[meta_id] = []\n",
    "    popularity_tags[meta_id].extend(meta[\"tags\"])\n",
    "\n",
    "# for each meta_id remove duplicates\n",
    "for meta_id, tags in popularity_tags.items():\n",
    "    popularity_tags[meta_id] = list(set(tags))\n",
    "\n",
    "print(f\"total number of unique ids: {len(popularity_tags)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all unique popularity tags\n",
    "popularity_tags_list = list(set(tag for tags in popularity_tags.values() for tag in tags))\n",
    "print(f\"total number of unique popularity tags: {len(popularity_tags_list)}\")\n",
    "# get the top 100 popularity tags\n",
    "# Count occurrences of each tag\n",
    "tag_counts = {}\n",
    "for tags in popularity_tags.values():\n",
    "    for tag in tags:\n",
    "        if tag in tag_counts:\n",
    "            tag_counts[tag] += 1\n",
    "        else:\n",
    "            tag_counts[tag] = 1\n",
    "# Sort by count\n",
    "popularity_tags_list = sorted(popularity_tags_list, key=lambda x: tag_counts.get(x, 0), reverse=True)[:1000]\n",
    "for tag in popularity_tags_list:\n",
    "    print(f\"{tag_counts[tag]}: {tag}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the ear scores \n",
    "discogs_path = \"/home/christian/code/christian/metadata/ear/discogs_subset_ear_scores.csv\"\n",
    "genius_path = \"/home/christian/code/christian/metadata/ear/genius_ear_scores.csv\"\n",
    "imslp_path = \"/home/christian/code/christian/metadata/ear/imslp_ear_scores.csv\"\n",
    "\n",
    "ear_scores = {}\n",
    "for path in [discogs_path, genius_path, imslp_path]:\n",
    "    # read the csv file\n",
    "    df = pl.read_csv(path)\n",
    "    # add the ear score to the ear_scores dictionary indexed by id\n",
    "    dict_df = df.to_dicts()\n",
    "    for row in dict_df:\n",
    "        ear_scores[row[\"id\"]] = [f\"\"\"quality: {int(row[\"mean_score\"])}\"\"\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# /home/christian/code/christian/metadata/organized/titles\n",
    "\n",
    "# load the youtube titles terms\n",
    "dicsogs_subset_title_terms_filepath = \"/home/christian/code/christian/metadata/terms/discogs_subset_title_terms_map.json\"\n",
    "with open(dicsogs_subset_title_terms_filepath, \"r\") as f:\n",
    "    dicsogs_subset_title_terms = json.load(f)\n",
    "\n",
    "genius_title_terms_filepath = \"/home/christian/code/christian/metadata/terms/genius_title_terms_map.json\"\n",
    "with open(genius_title_terms_filepath, \"r\") as f:\n",
    "    genius_title_terms = json.load(f)\n",
    "\n",
    "discogs_title_terms_filepath = \"/home/christian/code/christian/metadata/terms/discogs_title_terms_map.json\"\n",
    "with open(discogs_title_terms_filepath, \"r\") as f:\n",
    "    discogs_title_terms = json.load(f)\n",
    "\n",
    "# merge the terms maps\n",
    "terms_map = {**dicsogs_subset_title_terms, **genius_title_terms, **discogs_title_terms}\n",
    "print(len(terms_map))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the audio production metas\n",
    "combined_audio_features = pl.read_csv(\"/home/christian/code/christian/metadata/v4/combined_audio_production_features_v2.csv\")\n",
    "\n",
    "features = combined_audio_features.columns\n",
    "features = [f for f in features if \"normalized\" in f]\n",
    "# remove silence_percentage\n",
    "features = [f for f in features if \"silence_percentage\" not in f]\n",
    "print(features)\n",
    "\n",
    "audio_features = {}\n",
    "\n",
    "# iterate over the rows of the combined_audio_features\n",
    "for row in tqdm(combined_audio_features.iter_rows(named=True)):\n",
    "    # create tags from the normalized features\n",
    "    tags = []\n",
    "    for feature in features:\n",
    "        # convert the feature value to a tag\n",
    "        feature_value = row[feature]\n",
    "        if feature_value is None:\n",
    "            continue\n",
    "        feature_string = f\"{feature.replace('_normalized', '')}: {feature_value:.1f}\"\n",
    "        tags.append(feature_string)\n",
    "    audio_features[row[\"id\"]] = tags\n",
    "\n",
    "# load the metas_val\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get all unique audio feature tags\n",
    "audio_feature_tags = list(set(tag for tags in audio_features.values() for tag in tags))\n",
    "print(f\"Total number of unique audio_features tags: {len(audio_feature_tags)}\")\n",
    "\n",
    "# Count occurrences of each popularity tag\n",
    "tag_counts = {}\n",
    "for tags in audio_features.values():\n",
    "    for tag in tags:\n",
    "        if tag in tag_counts:\n",
    "            tag_counts[tag] += 1\n",
    "        else:\n",
    "            tag_counts[tag] = 1\n",
    "\n",
    "# Create a list of all popularity tags\n",
    "audio_features_list = list(tag_counts.keys())\n",
    "\n",
    "# Sort by count and get top 1000\n",
    "audio_features_list = sorted(audio_features_list, key=lambda x: tag_counts.get(x, 0), reverse=True)[:1000]\n",
    "\n",
    "# Print the top popularity tags with their counts\n",
    "for tag in audio_features_list:\n",
    "    print(f\"{tag_counts[tag]}: {tag}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# as a final step, lets merge all the tags into one metadata file for convenience\n",
    "combined_tags = {}\n",
    "# Merge all tag dictionaries, combining lists for the same ids\n",
    "for tag_dict in [gpt_tags, ear_scores, popularity_tags, audio_features, terms_map]:\n",
    "    for id, tags in tqdm(tag_dict.items()):\n",
    "        if id not in combined_tags:\n",
    "            combined_tags[id] = []\n",
    "        if isinstance(tags, list):\n",
    "            combined_tags[id].extend(tags)\n",
    "        else:\n",
    "            combined_tags[id].append(tags)\n",
    "\n",
    "# Count the total number of tags across all items\n",
    "total_tags = sum(len(tags) for tags in combined_tags.values())\n",
    "print(f\"Number of items with tags: {len(combined_tags)}\")\n",
    "print(f\"Total number of tags across all items: {total_tags}\")\n",
    "print(f\"Average tags per item: {total_tags / len(combined_tags):.2f}\")\n",
    "\n",
    "# save the combined tags\n",
    "#with open(\"/home/christian/code/christian/metadata/combined_tags_apr30.json\", \"w\") as f:\n",
    "#    json.dump(combined_tags, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = \"/app/suno/data/auk_v0/\"\n",
    "subset = \"tr\"\n",
    "metas_val_path = f\"metas_v1_{subset}.jsonl\"\n",
    "metas_val = read_jsonl(os.path.join(base_dir, metas_val_path))\n",
    "print(len(metas_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#base_dir = \"/app/suno/data/diffusion_v5/v0\"\n",
    "#metas_val_path = \"metas_val_aligned_v2.jsonl\"\n",
    "subset = \"tr\"\n",
    "base_dir = \"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz\"\n",
    "metas_val_path = f\"metas_{subset}.jsonl\"\n",
    "metas_val = pl.read_ndjson(os.path.join(base_dir, metas_val_path))\n",
    "print(len(metas_val))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tag_count = 0\n",
    "ear_count = 0\n",
    "chart_count = 0\n",
    "audio_count = 0\n",
    "terms_count = 0\n",
    "total_additional_tags = 0\n",
    "\n",
    "# new metas\n",
    "new_metas = []\n",
    "for meta_idx, meta in enumerate(tqdm(metas_val)):\n",
    "    new_meta = meta.copy()\n",
    "    # Initialize tags if it's None\n",
    "\n",
    "    # check if clip_id is in gpt_tags\n",
    "    if meta[\"id\"] in gpt_tags:\n",
    "        tag_count += 1\n",
    "        total_additional_tags += len(gpt_tags[meta[\"id\"]])\n",
    "        if new_meta.get(\"tags\") is None:\n",
    "            new_meta[\"tags\"] = []\n",
    "        new_meta[\"tags\"].extend(gpt_tags[meta[\"id\"]])\n",
    "    if meta[\"id\"] in ear_scores:\n",
    "        ear_count += 1\n",
    "        total_additional_tags += 1\n",
    "        if new_meta.get(\"tags\") is None:\n",
    "            new_meta[\"tags\"] = []\n",
    "        new_meta[\"tags\"].extend(ear_scores[meta[\"id\"]])\n",
    "    if meta[\"id\"] in popularity_tags:\n",
    "        chart_count += 1\n",
    "        total_additional_tags += len(popularity_tags[meta[\"id\"]])\n",
    "        if new_meta.get(\"tags\") is None:\n",
    "            new_meta[\"tags\"] = []\n",
    "        new_meta[\"tags\"].extend(popularity_tags[meta[\"id\"]])\n",
    "    if meta[\"id\"] in audio_features:\n",
    "        audio_count += 1\n",
    "        total_additional_tags += len(audio_features[meta[\"id\"]])\n",
    "        if new_meta.get(\"tags\") is None:\n",
    "            new_meta[\"tags\"] = []\n",
    "        new_meta[\"tags\"].extend(audio_features[meta[\"id\"]])\n",
    "    if meta[\"id\"] in terms_map:\n",
    "        terms_count += 1\n",
    "        total_additional_tags += len(terms_map[meta[\"id\"]])\n",
    "        if new_meta.get(\"tags\") is None:\n",
    "            new_meta[\"tags\"] = []\n",
    "        new_meta[\"tags\"].extend(terms_map[meta[\"id\"]])\n",
    "\n",
    "    new_metas.append(new_meta)\n",
    "\n",
    "print(tag_count, len(metas_val))\n",
    "print(ear_count, len(metas_val))\n",
    "print(chart_count, len(metas_val))\n",
    "print(audio_count, len(metas_val))\n",
    "print(terms_count, len(metas_val))\n",
    "print(total_additional_tags, len(metas_val))\n",
    "print(len(new_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for meta in new_metas:\n",
    "    if \"playlist_ids\" in meta:\n",
    "        if meta[\"playlist_ids\"] is not None:\n",
    "            print(\"found playlist_ids \", meta[\"id\"], meta[\"playlist_ids\"])\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(new_metas, f\"/mnt/localdisk/tmp_cjs/metas_v2_{subset}.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_metas[3307]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_metas[3307]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_metas[3307]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_quality = 100\n",
    "# go through metas and find one with lowest quality\n",
    "for meta in metas_val.iter_rows(named=True):\n",
    "    if meta[\"dataset\"] == \"imslp\":\n",
    "        continue\n",
    "    if meta[\"id\"] in ear_scores:\n",
    "        quality_score = ear_scores[meta[\"id\"]][0].split(\":\")[1].strip()\n",
    "        if int(quality_score) < min_quality:\n",
    "            min_quality = int(quality_score)\n",
    "            min_quality_id = meta[\"id\"]\n",
    "            min_quality_meta = meta\n",
    "\n",
    "print(min_quality, min_quality_meta[\"id\"], min_quality_meta[\"tags\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}