{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = \"discogs_subset\"\n",
    "\n",
    "base_dir = f\"/home/christian/ear_scores/{dataset_name}\"\n",
    "os.makedirs(base_dir, exist_ok=True)\n",
    "# find all json files in the base_dir\n",
    "json_files = [f for f in os.listdir(base_dir) if f.endswith('.json')]\n",
    "print(len(json_files))\n",
    "json_filepaths = [os.path.join(base_dir, f) for f in json_files]\n",
    "\n",
    "# read the json files and store into one dict\n",
    "data = {}\n",
    "for f in tqdm(json_filepaths):\n",
    "    data.update(json.load(open(f)))\n",
    "\n",
    "# create a dataframe with just the track id and mean score\n",
    "rows = []\n",
    "for track_id, track_data in data.items():\n",
    "    row = {\n",
    "        'id': track_id,\n",
    "        'mean_score': track_data['mean_score']\n",
    "    }\n",
    "    rows.append(row)\n",
    "\n",
    "df = pd.DataFrame(rows)\n",
    "# Set id as index\n",
    "df.set_index('id', inplace=True)\n",
    "\n",
    "print(df.head())\n",
    "print(df.describe())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save this final json file to disk\n",
    "output_path = f\"/home/christian/code/christian/metadata/ear/{dataset_name}_ear_scores.json\"\n",
    "with open(output_path, 'w') as f:\n",
    "    json.dump(data, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the dataframe to disk \n",
    "output_path = f\"/home/christian/code/christian/metadata/ear/{dataset_name}_ear_scores.csv\"\n",
    "df.to_csv(output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = \"genius\"\n",
    "# read the csv file\n",
    "output_path = f\"/home/christian/code/christian/metadata/ear/{dataset_name}_ear_scores.csv\"\n",
    "\n",
    "df = pd.read_csv(output_path)\n",
    "# make a histogram of the mean scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a histogram of the mean scores\n",
    "plt.hist(df['mean_score'], bins=250, zorder=2, alpha=0.8)\n",
    "# compute 95% percentile of the mean scores and plot a vertical line\n",
    "print(df['mean_score'].quantile(0.9))\n",
    "plt.axvline(df['mean_score'].quantile(0.9), color='red', linestyle='--', zorder=3)\n",
    "# and 5%\n",
    "plt.axvline(df['mean_score'].quantile(0.10), color='red', linestyle='--')\n",
    "# count the number of tracks below the 5% threshold\n",
    "print(len(df[df['mean_score'] < df['mean_score'].quantile(0.1)]))\n",
    "#plt.yscale('log')\n",
    "plt.grid(c=\"lightgray\")\n",
    "plt.xlabel(\"Mean Score\")\n",
    "plt.ylabel(\"Number of Tracks\")\n",
    "plt.title(\"Distribution of Mean Ear scores\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.s3 import read_from_s3\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "metas_s3_filepath = f\"s3://suno-data/datasets/bundles/v4/{dataset_name}/metas_v0.jsonl\"\n",
    "metas = read_from_s3(metas_s3_filepath, read_f=read_jsonl)\n",
    "print(len(metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas_map = {meta['id']: meta for meta in metas}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort dataframe by mean_score from low to high\n",
    "df = df.sort_values(by='mean_score', ascending=False)\n",
    "print(df.head(10))\n",
    "\n",
    "# get tracks that have a score below 15\n",
    "#low_score_tracks = df[df['mean_score'] < 15]\n",
    "#print(low_score_tracks.head(10))\n",
    "\n",
    "# sort dataframe by mean_score from low to high\n",
    "df = df.sort_values(by='mean_score', ascending=True)\n",
    "print(df.head(10))\n",
    "\n",
    "\n",
    "\n",
    "# random 5 tracks\n",
    "#print(df.sample(10))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# given an s3 filepath, download the audio and print the local path\n",
    "def download_audio(s3_filepath, local_filepath):\n",
    "    # download the audio from s3\n",
    "    os.system(f\"aws s3 cp {s3_filepath} {local_filepath}\")\n",
    "\n",
    "track_id = \"KjekHwmsvvY\"\n",
    "# get the s3 filepath\n",
    "s3_filepath = metas_map[track_id]['s3_filepath']\n",
    "print(metas_map[track_id])\n",
    "local_filepath = f\"/mnt/localdisk/tmp/audio/{track_id}.mp3\"\n",
    "download_audio(s3_filepath, local_filepath)\n",
    "print(f\"Downloaded {s3_filepath} to {local_filepath}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load old metas to map back to old ids\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "old_genius_metas = read_jsonl(\"/home/christian/code/christian/metadata/genius_hq_metas.jsonl\")\n",
    "genius_id_map = {meta[\"original_id\"]: meta[\"id\"] for meta in old_genius_metas}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the raw genius metas to get the view counts \n",
    "raw_genius_metas = read_jsonl(\"/app/suno/tmp/raw_genius_metas.jsonl\")\n",
    "raw_genius_metas_map = {meta[\"id\"]: meta for meta in raw_genius_metas}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "# we can also cross section with the old school audio features\n",
    "audio_prod_path = \"/home/christian/code/christian/metadata/genius_hq_audio_production_features_v2.csv\"\n",
    "audio_prod = pd.read_csv(audio_prod_path)\n",
    "audio_features_map = {m[\"id\"]: m for m in audio_prod.to_dict(orient=\"records\")}\n",
    "\n",
    "feature_bounds = {\n",
    "    \"loudness\": [-24, -2],\n",
    "    \"spectral_centroid\": [1500, 5000],\n",
    "    \"spectral_flatness\": [0.02, 0.3],\n",
    "    #\"crest_factor\": [1.0, 3],\n",
    "    #\"bass\" : [0.1, 0.5],\n",
    "    #\"mid\" : [0.4, 1.0],\n",
    "    #\"high\" : [0.15, 1.25],\n",
    "    \"stereo_width\" : [0.15, 0.4]\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can save out a json file with the track id at our given cutoff\n",
    "#output_filepath = f\"/home/christian/code/christian/metadata/ear/genius_t1.json\" # cutoff below 20 and above 26\n",
    "\n",
    "data_cut = \"t7\"\n",
    "output_filepath = f\"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz/info_tr_{data_cut}.json\" # cutoff below 20 and above 26\n",
    "\n",
    "print(f\"number of tracks before cutoff: {len(df)}\")\n",
    "#passed_tracks = df[(df['mean_score'] > 20) & (df['mean_score'] < 25)] # cut on both tails\n",
    "passed_tracks = df[(df['mean_score'] > 22) & (df['mean_score'] < 28)] # allow more in the top tail\n",
    "print(f\"number of tracks after ear cutoff: {len(passed_tracks)}\")\n",
    "\n",
    "passed_tracks = passed_tracks[\"id\"].tolist()\n",
    "\n",
    "# filter based on views\n",
    "passed_tracks = [id for id in passed_tracks if raw_genius_metas_map[id][\"views\"] > 10_000]\n",
    "print(f\"number of tracks after views filter: {len(passed_tracks)}\")\n",
    "\n",
    "# covert passed tracks to the old ids\n",
    "passed_tracks_old_ids = [genius_id_map[id] for id in passed_tracks]\n",
    "\n",
    "new_passed_tracks_old_ids = []\n",
    "new_passed_tracks = []\n",
    "# now we can filter the audio features\n",
    "for track_id in tqdm(passed_tracks):\n",
    "    audio_features = audio_features_map[genius_id_map[track_id]]\n",
    "    is_within_bounds = True\n",
    "    for feature in feature_bounds:\n",
    "        if audio_features[feature] < feature_bounds[feature][0] or audio_features[feature] > feature_bounds[feature][1]:\n",
    "            is_within_bounds = False\n",
    "            break\n",
    "    \n",
    "    if is_within_bounds:\n",
    "        new_passed_tracks.append(track_id)\n",
    "        new_passed_tracks_old_ids.append(genius_id_map[track_id])\n",
    "\n",
    "print(f\"number of tracks after filtering: {len(new_passed_tracks)}\")\n",
    "# create a dict with the dataset name\n",
    "output_dict = {\n",
    "    \"diffusion_mix_fix\": new_passed_tracks_old_ids\n",
    "}\n",
    "print(output_dict.keys())\n",
    "# save out the passed tracks\n",
    "with open(output_filepath, 'w') as f:\n",
    "   json.dump(output_dict, f)\n",
    "\n",
    "print(len(output_dict[\"diffusion_mix_fix\"]))\n",
    "print(output_dict[\"diffusion_mix_fix\"][:10])\n",
    "\n",
    "\n",
    "output_filepath = f\"/home/christian/code/christian/metadata/ear/{dataset_name}_{data_cut}.json\"\n",
    "output_dict = {\n",
    "    f\"{dataset_name}\": new_passed_tracks\n",
    "}\n",
    "print(output_dict.keys())\n",
    "# save out the passed tracks\n",
    "with open(output_filepath, 'w') as f:\n",
    "   json.dump(output_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dir = \"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz\"\n",
    "info_filename = \"info_tr_t2.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load info.json if it exists to filter metas\n",
    "filtered_indices = None\n",
    "if info_filename is not None and os.path.exists(\n",
    "    os.path.join(dataset_dir, info_filename)\n",
    "):\n",
    "    with open(os.path.join(dataset_dir, self.info_filename), \"r\") as f:\n",
    "        info_data = json.load(f)\n",
    "\n",
    "    # Create a set of all allowed IDs across all dataset keys\n",
    "    allowed_ids = defaultdict(set)\n",
    "    for dataset_name, id_list in info_data.items():\n",
    "        allowed_ids[dataset_name].update(id_list)\n",
    "\n",
    "    # Create a mapping from original indices to filtered indices\n",
    "    self.filtered_indices = []\n",
    "    for i, meta in enumerate(self.metas):\n",
    "        dataset = meta.get(\"dataset\")\n",
    "        if dataset == \"discogs_subset\":\n",
    "            dataset = \"discogs\"\n",
    "        if dataset in allowed_ids:\n",
    "            if meta.get(\"id\") in allowed_ids[dataset]:\n",
    "                self.filtered_indices.append(i)\n",
    "\n",
    "    assert len(self.filtered_indices) > 0, \"No samples left after filtering\"\n",
    "\n",
    "    print_with_time_master(\n",
    "        f\"Filtered dataset using {self.info_filename}: {len(self.filtered_indices)} / {len(self.metas)} samples kept\"\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}