{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import polars as pl\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for this to work, we will first have to load the old data info\n",
    "# from this we can the list of ids for each "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# also load the json info that we used \n",
    "import pandas as pd\n",
    "\n",
    "output_path = \"/home/christian/code/christian/metadata/ear/genius_ear_scores.csv\"\n",
    "df = pd.read_csv(output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "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": [
    "# lets add a new column to the df for \"original_id\" using the genius_id_map\n",
    "# Create a mapping from id to original_id (reverse of genius_id_map)\n",
    "reverse_genius_id_map = {v: k for k, v in genius_id_map.items()}\n",
    "\n",
    "# Add the original_id column to the dataframe\n",
    "df['original_id'] = df['id'].map(reverse_genius_id_map)\n",
    "\n",
    "# Check if any mappings are missing\n",
    "missing_ids = df[df['original_id'].isna()]['id'].tolist()\n",
    "if missing_ids:\n",
    "    print(f\"Warning: {len(missing_ids)} IDs could not be mapped to original_ids\")\n",
    "    print(f\"First few missing IDs: {missing_ids[:5]}\")\n",
    "else:\n",
    "    print(\"All IDs successfully mapped to original_ids\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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, -8],\n",
    "    \"spectral_centroid\": [2000, 4500],\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": [
    "from tqdm import tqdm\n",
    "\n",
    "print(f\"number of tracks before cutoff: {len(df)}\")\n",
    "passed_tracks = df[(df['mean_score'] > 20) & (df['mean_score'] < 25)]\n",
    "print(f\"number of tracks after cutoff: {len(passed_tracks)}\")\n",
    "\n",
    "# covert passed tracks to the old ids\n",
    "passed_tracks = passed_tracks[\"id\"].tolist()\n",
    "passed_tracks = [genius_id_map[id] for id in passed_tracks]\n",
    "\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[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",
    "\n",
    "print(f\"number of tracks after filtering: {len(new_passed_tracks)}\")\n",
    "# create a dict with the dataset name\n",
    "output_dict = {\n",
    "    \"genius\": 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)\n",
    "\n",
    "print(len(output_dict[\"genius\"]))\n",
    "print(output_dict[\"genius\"][:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preload the source metas\n",
    "source_metas = {}\n",
    "for subset in [\"val\", \"tr\"]:\n",
    "    source_metas[subset] = pl.read_ndjson(f\"/app/suno/data/chirp_v5_ft/v2/metas_{subset}.jsonl\")\n",
    "    print(f\"subset {subset}: {len(source_metas[subset])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "for subset in [\"tr\"]:\n",
    "\n",
    "    source_metas_subset = source_metas[subset]\n",
    "\n",
    "    with open(f\"/app/suno/data/chirp_v5_ft/v2/info_{subset}.json\", \"r\") as f:\n",
    "        original_info = json.load(f)\n",
    "    \n",
    "    for meta_idx, meta in enumerate(tqdm(source_metas_subset)):\n",
    "        dataset = meta[\"dataset\"]\n",
    "        task = meta[\"task\"]\n",
    "\n",
    "        if \"genius\" in dataset: # apply our filter\n",
    "            "
   ]
  }
 ],
 "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
}