{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.utils.s3 import read_from_s3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_metas_filepath = \"/home/christian/code/christian/metadata/genius_hq_metas.jsonl\"\n",
    "base_metas = read_jsonl(base_metas_filepath)\n",
    "print(f\"Loaded {len(base_metas)} base metas\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load alignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load audio production metas\n",
    "audio_production_metas_filepath = \"/home/christian/code/christian/metadata/genius_hq_metas_audio_production.jsonl\"\n",
    "audio_production_metas = read_jsonl(audio_production_metas_filepath)\n",
    "print(f\"Loaded {len(audio_production_metas)} audio production metas\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_production_metas_map = {meta[\"id\"]: meta for meta in audio_production_metas}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_metas =[]\n",
    "from tqdm import tqdm\n",
    "for meta in tqdm(base_metas):\n",
    "    if meta[\"id\"] in audio_production_metas_map:\n",
    "        new_meta = meta.copy()\n",
    "        new_meta[\"features\"] = audio_production_metas_map[meta[\"id\"]][\"features\"]\n",
    "        merged_metas.append(new_meta)\n",
    "    else:\n",
    "        print(f\"Meta {meta['id']} not found in audio production metas\")\n",
    "\n",
    "print(f\"Merged {len(merged_metas)} metas\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create histogram of audio production metas\n",
    "\n",
    "# Extract spectral centroids for different categories\n",
    "all_centroids = [float(entry['features']['spectral_centroid']) for entry in merged_metas]\n",
    "en_centroids = [float(entry['features']['spectral_centroid']) for entry in merged_metas \n",
    "                if entry['lang'] == 'en']\n",
    "other_centroids = [float(entry['features']['spectral_centroid']) for entry in merged_metas \n",
    "                    if entry['lang'] != 'en']\n",
    "\n",
    "# remove nan values\n",
    "all_centroids = [val for val in all_centroids if np.isfinite(val)]\n",
    "en_centroids = [val for val in en_centroids if np.isfinite(val)]\n",
    "other_centroids = [val for val in other_centroids if np.isfinite(val)]\n",
    "\n",
    "print(f\"Found {len(all_centroids)} centroids\")\n",
    "print(f\"Found {len(en_centroids)} en lang centroids\")\n",
    "print(f\"Found {len(other_centroids)} other lang centroids\")\n",
    "\n",
    "# Create single figure\n",
    "plt.figure(figsize=(12, 6))\n",
    "bins = 1000\n",
    "\n",
    "# Plot normalized histograms (density=True makes the area under each histogram equal to 1)\n",
    "plt.hist(all_centroids, bins=bins, density=True, color='blue', alpha=0.3, label='All Entries')\n",
    "plt.hist(en_centroids, bins=bins, density=True, color='green', alpha=0.3, label='English Entries')\n",
    "plt.hist(other_centroids, bins=bins, density=True, color='red', alpha=0.3, label='Non-English Entries')\n",
    "\n",
    "plt.title('Normalized Spectral Centroid Distributions', fontsize=16)\n",
    "plt.xlabel('Spectral Centroid')\n",
    "plt.ylabel('Density')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.legend()\n",
    "\n",
    "# Add statistics text\n",
    "stats_text = (\n",
    "    f'Total entries: {len(all_centroids)}\\n'\n",
    "    f'English entries: {len(en_centroids)}\\n'\n",
    "    f'Other entries: {len(other_centroids)}\\n\\n'\n",
    "    f'Mean (All): {np.mean(all_centroids):.2f}\\n'\n",
    "    f'Mean (EN): {np.mean(en_centroids):.2f}\\n'\n",
    "    f'Mean (Other): {np.mean(other_centroids):.2f}\\n\\n'\n",
    "    f'Median (All): {np.median(all_centroids):.2f}\\n'\n",
    "    f'Median (EN): {np.median(en_centroids):.2f}\\n'\n",
    "    f'Median (Other): {np.median(other_centroids):.2f}'\n",
    ")\n",
    "plt.text(0.02, 0.98, stats_text, \n",
    "         transform=plt.gca().transAxes,\n",
    "         verticalalignment='top',\n",
    "         fontsize=10,\n",
    "         bbox=dict(facecolor='white', alpha=0.8))\n",
    "\n",
    "# Add vertical lines for means\n",
    "for data, color, name in [(all_centroids, 'blue', 'All'), \n",
    "                         (en_centroids, 'green', 'English'), \n",
    "                         (other_centroids, 'red', 'Other')]:\n",
    "    mean_val = np.mean(data)\n",
    "    plt.axvline(mean_val, color=color, linestyle='--', alpha=0.5)\n",
    "    plt.text(mean_val, plt.gca().get_ylim()[1], f'{name} mean', \n",
    "             rotation=90, va='top', ha='right', color=color)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter out metas for features that out of distribution in 1% and 99%\n",
    "features = [\"spectral_centroid\", \"loudness_factor\"]\n",
    "for feature in features:\n",
    "    feature_values = [float(meta[\"features\"][feature]) for meta in audio_production_metas]\n",
    "    feature_values = [val for val in feature_values if np.isfinite(val)]\n",
    "    q1 = np.quantile(feature_values, 0.01)\n",
    "    q99 = np.quantile(feature_values, 0.99)\n",
    "    print(f\"{feature}: q1 = {q1}, q99 = {q99}\")\n",
    "\n",
    "    # how many metas are out of distribution?\n",
    "    out_of_dist_metas = [meta for meta in audio_production_metas if float(meta[\"features\"][feature]) < q1 or float(meta[\"features\"][feature]) > q99]\n",
    "    print(f\"Found {len(out_of_dist_metas)} metas out of distribution for {feature}.\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter out metas with less than 100 youtube views\n",
    "filtered_base_metas1 = [meta for meta in base_metas if meta[\"youtube_views\"] >= 100]\n",
    "print(f\"Filtered to {len(filtered_base_metas1)} metas with at least 100 youtube views. Removed {len(base_metas) - len(filtered_base_metas1)} metas.\")\n",
    "\n",
    "# filter out metas with duration less than 1 minute and greater than 8 minutes\n",
    "filtered_base_metas2 = [meta for meta in filtered_base_metas1 if meta[\"duration_s\"] >= 60 and meta[\"duration_s\"] <= 480]\n",
    "print(f\"Filtered to {len(filtered_base_metas2)} metas with duration between 1 minute and 8 minutes. Removed {len(filtered_base_metas1) - len(filtered_base_metas2)} metas.\")\n",
    "\n",
    "# filter out metas with more than 6144 characters in lyrics or less than 50 characters in lyrics\n",
    "filtered_base_metas3 = [meta for meta in filtered_base_metas2 if len(meta[\"lyrics\"]) <= 6144 and len(meta[\"lyrics\"]) >= 50]\n",
    "print(f\"Filtered to {len(filtered_base_metas3)} metas with lyrics of at most 6144 characters and at least 50 characters. Removed {len(filtered_base_metas2) - len(filtered_base_metas3)} metas.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count greek lang in metas\n",
    "greek_metas = [meta for meta in filtered_base_metas3 if meta[\"lang\"] == \"el\"]\n",
    "print(f\"Found {len(greek_metas)} metas in Greek.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check for number of metas with \"metal\" tag\n",
    "tag_to_check = \"epic\"\n",
    "\n",
    "metal_metas = []\n",
    "for meta in filtered_base_metas3:\n",
    "    tags = meta[\"tags_text\"]\n",
    "    tags = [tag.strip().lower() for tag in tags]\n",
    "    tags = [tag.replace(\"genius\", \"\").strip() for tag in tags]\n",
    "    if tag_to_check in tags:\n",
    "        metal_metas.append(meta)\n",
    "print(f\"Found {len(metal_metas)} metas with '{tag_to_check}' tag.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_base_metas3[100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion_metas_filepath = \"/home/christian/code/christian/metadata/diffusion_mix_v2/metas.jsonl\"\n",
    "diffusion_metas = read_jsonl(diffusion_metas_filepath)\n",
    "print(f\"Loaded {len(diffusion_metas)} diffusion metas\")\n",
    "\n",
    "# count the number of metas with lyrics\n",
    "diffusion_metas_with_lyrics = [meta for meta in diffusion_metas if \"text\" in meta]\n",
    "print(f\"Found {len(diffusion_metas_with_lyrics)} diffusion metas with lyrics\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "alignments = read_jsonl(\n",
    "    \"/home/christian/code/christian/metadata/ytm_alignments_v11.jsonl\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "alignments_map = {}\n",
    "for meta_id, alignment in alignments:\n",
    "    alignments_map[meta_id] = alignment\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "counter = 0\n",
    "for diffusion_meta in diffusion_metas:\n",
    "    if diffusion_meta[\"id\"] in alignments_map:\n",
    "        #diffusion_meta[\"alignments\"] = alignments_map[diffusion_meta[\"id\"]]\n",
    "        counter += 1\n",
    "print(f\"Found {counter} metas with aligned lyrics\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check current metas for langauge and lyrics\n",
    "\n",
    "base_dir = \"/app/suno/data/diffusion_mix/vae_25hz_30s\"\n",
    "train_metas_filepath = \"metas_context_aligned_tr.jsonl\"\n",
    "\n",
    "train_metas = read_jsonl(os.path.join(base_dir, train_metas_filepath))\n",
    "print(f\"Loaded {len(train_metas):,} train metas\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_metas[100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get number that have lyrics (\"Text\")\n",
    "lyrics_metas = [meta for meta in train_metas if \"text\" in meta]\n",
    "print(f\"Found {len(lyrics_metas):,} metas with lyrics ({len(lyrics_metas) / len(train_metas) * 100:.2f}%)\")\n",
    "\n",
    "# get number that have alignments\n",
    "alignments_metas = [meta for meta in train_metas if \"text_aligned\" in meta]\n",
    "print(f\"Found {len(alignments_metas):,} metas with alignments ({len(alignments_metas) / len(train_metas) * 100:.2f}%)\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count languages\n",
    "langs = [meta.get(\"text_lang\", None) for meta in train_metas]\n",
    "print(f\"Found {len(set(langs))} languages: {set(langs)}\")\n",
    "\n",
    "# count occurance of each language\n",
    "lang_counts = {lang: langs.count(lang) for lang in set(langs)}\n",
    "print(lang_counts)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort lang counts by value\n",
    "sorted_lang_counts = sorted(lang_counts.items(), key=lambda x: x[1], reverse=True)\n",
    "print(sorted_lang_counts)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check greek\n",
    "greek_count = lang_counts.get(\"el\", 0)\n",
    "print(f\"Found {greek_count} greek metas\")\n",
    "# conver to hours, each file is 30s\n",
    "greek_hours = greek_count * 30 / 3600\n",
    "print(f\"Found {greek_hours:.2f} hours of greek metas\")\n",
    "print()\n",
    "\n",
    "# check arabic\n",
    "arabic_count = lang_counts.get(\"ar\", 0)\n",
    "print(f\"Found {arabic_count} arabic metas\")\n",
    "# conver to hours, each file is 30s\n",
    "arabic_hours = arabic_count * 30 / 3600\n",
    "print(f\"Found {arabic_hours:.2f} hours of arabic metas\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for lang, count in sorted_lang_counts:\n",
    "    # conver to hours, each file is 30s\n",
    "    hours = int(round(count * 30 / 3600, 2))\n",
    "    print(f\"{lang}: {hours:,} hours\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count how many languages have less than 1 hour of data\n",
    "less_than_1_hour = [count for count in lang_counts.values() if count * 30 / 3600 < 1]\n",
    "print(f\"Found {len(less_than_1_hour)} languages with less than 1 hour of data.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count how many are english\n",
    "english_count = lang_counts.get(\"en\", 0)\n",
    "print(f\"{english_count:,} english metas\")\n",
    "# conver to hours, each file is 30s\n",
    "english_hours = english_count * 30 / 3600\n",
    "print(f\"{english_hours:.2f} hours of english metas\")\n",
    "\n",
    "\n",
    "# count how many are non-english (anything else except None)\n",
    "non_english_count = 0\n",
    "for lang in lang_counts:\n",
    "    if lang != \"en\" and lang is not None:\n",
    "        non_english_count += lang_counts[lang]\n",
    "print(f\"{non_english_count:,} non-english metas\")\n",
    "# conver to hours, each file is 30s\n",
    "non_english_hours = non_english_count * 30 / 3600\n",
    "print(f\"{non_english_hours:.2f} hours of non-english metas\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate total hours\n",
    "total_hours = len(train_metas) * 30 / 3600\n",
    "total_hours_rounded = int(round(total_hours, 2))\n",
    "print(f\"{total_hours_rounded:,} hours of total data\")\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
}