{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import tempfile\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3, upload_s3_files\n",
    "from suno_utils.audio import Audio\n",
    "import emoji"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_metas = read_from_s3(\"s3://suno-data/datasets/harvest/splice/splice_all_samples_data_cleaned.jsonl\", read_f=read_jsonl)\n",
    "df_raw = pd.DataFrame.from_dict(raw_metas)\n",
    "print(df_raw.columns)\n",
    "df_raw.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def strip_text(text):\n",
    "    if text is None:\n",
    "        return text\n",
    "    text = text.lower()\n",
    "    text = re.sub(r'[^a-zA-Z0-9\\s]', '', text)  # Remove non-alphanumeric except spaces\n",
    "    text = re.sub(r'\\s+', ' ', text)  # Replace multiple spaces with single space\n",
    "    return text.strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "### Clean Splice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean = df_raw.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean[\"duration_s\"] = df_clean[\"duration\"].apply(lambda x: x / 1000.0 if isinstance(x, int) else None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_bpm(bpm_str):\n",
    "    if pd.isna(bpm_str) or bpm_str == 'None':\n",
    "        return None\n",
    "    \n",
    "    try:\n",
    "        # Convert to float first\n",
    "        bpm_float = float(bpm_str)\n",
    "        \n",
    "        # Round to 1 decimal place\n",
    "        bpm_rounded = round(bpm_float, 1)\n",
    "        \n",
    "        # If it's a whole number, return as integer string\n",
    "        if bpm_rounded.is_integer():\n",
    "            return str(int(bpm_rounded))\n",
    "        else:\n",
    "            return f\"{bpm_rounded:.1f}\"\n",
    "    except (ValueError, TypeError):\n",
    "        return None\n",
    "\n",
    "# Apply the normalization\n",
    "df_clean['bpm'] = df_clean['bpm'].apply(normalize_bpm)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_12_keys = ['c', 'f', 'a', 'e', 'd', 'g', 'b', 'f#', 'd#', 'g#', 'a#', 'c#']\n",
    "\n",
    "def filter_key(key_value):\n",
    "    if pd.isna(key_value) or key_value is None:\n",
    "        return None\n",
    "    \n",
    "    # Check if the key is in the top 12\n",
    "    if key_value in top_12_keys:\n",
    "        return key_value\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "# Apply the filter\n",
    "df_clean['key'] = df_clean['key'].apply(filter_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_tags(tags_data):\n",
    "    genre_names = []\n",
    "    for tag in tags_data:\n",
    "        tag_label = tag.get('label', None)\n",
    "        if tag_label is not None and isinstance(tag_label, str) and len(tag_label) > 0:\n",
    "            genre_names.append(tag_label)\n",
    "    if len(genre_names) == 0:\n",
    "        return None\n",
    "    return genre_names\n",
    "\n",
    "df_clean['tags'] = df_clean['tags'].apply(parse_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean['parent_slug'] = df_clean['parent'].apply(lambda x: x.get('permalink_slug') if isinstance(x, dict) else None)\n",
    "df_clean['parent_base'] = df_clean['parent'].apply(lambda x: x.get('permalink_base_url') if isinstance(x, dict) else None)\n",
    "df_clean['s3_filepath'] = df_clean['uuid'].apply(lambda uuid: f\"s3://suno-data/datasets/harvest/splice/audio/{uuid}.mp3\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean = df_clean.drop(columns=[\"urls\", \"duration\", \"parent\"])\n",
    "df_clean.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean['duration_s'].clip(upper=250).hist(bins=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Filter Splice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter = df_clean.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(df_filter))\n",
    "df_filter = df_filter[df_filter['duration_s'] <= 15.0]\n",
    "print(len(df_filter))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fastest_dedupe_tags(tag_series):\n",
    "    # Convert to DataFrame with exploded tags\n",
    "    df_exploded = (tag_series.to_frame('tags')\n",
    "                   .explode('tags')\n",
    "                   .reset_index())\n",
    "    \n",
    "    # Quick filter and clean\n",
    "    mask = df_exploded['tags'].notna()\n",
    "    df_exploded = df_exploded[mask].copy()\n",
    "    \n",
    "    if df_exploded.empty:\n",
    "        return tag_series.apply(lambda x: [])\n",
    "    \n",
    "    df_exploded['tags'] = df_exploded['tags'].astype(str).str.strip()\n",
    "    df_exploded = df_exploded[df_exploded['tags'] != '']\n",
    "    \n",
    "    # Vectorized normalization\n",
    "    df_exploded['norm'] = df_exploded['tags'].apply(strip_text)\n",
    "    \n",
    "    # Remove empty normalized and deduplicate\n",
    "    df_exploded = df_exploded[df_exploded['norm'] != '']\n",
    "    df_exploded = df_exploded.drop_duplicates(['index', 'norm'], keep='first')\n",
    "    \n",
    "    # Regroup\n",
    "    return (df_exploded.groupby('index')['tags']\n",
    "            .apply(list)\n",
    "            .reindex(tag_series.index, fill_value=[]))\n",
    "\n",
    "# Apply\n",
    "df_filter['processed_tags'] = fastest_dedupe_tags(df_filter['tags'])\n",
    "df_filter = df_filter.drop(['tags'], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter.loc[df_filter['duration_s'] < 1.0, 'bpm'] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_values = df_filter['processed_tags'].explode().value_counts().head(25)\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "top_values.plot(kind='barh')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_values = df_filter['parent_slug'].value_counts().head(25)\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "top_values.plot(kind='barh')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(df_filter))\n",
    "df_filter.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "### Save Filtered Metas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_path = \"s3://suno-data/datasets/harvest/splice/sfx_metas_v0.jsonl\"\n",
    "with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=True) as temp_file:\n",
    "    df_filter.to_json(temp_file.name, orient='records', lines=True)\n",
    "    temp_filename = temp_file.name\n",
    "    results = upload_s3_files(from_local_filepaths=[temp_filename], to_s3_filepaths=[upload_path])\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "### Checking All the Filtered Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_metas = read_from_s3(\"s3://suno-data/datasets/harvest/splice/sfx_metas_v0.jsonl\", read_f=read_jsonl)\n",
    "print(f\"{len(raw_metas):,} tracks with {sum([m['duration_s'] for m in raw_metas])/60/60:,.1f}h total\")\n",
    "df_raw = pd.DataFrame.from_dict(raw_metas)\n",
    "print(df_raw.columns)\n",
    "df_raw.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_metas = read_from_s3(\"s3://suno-data/datasets/harvest/freesound/metas_v0.jsonl\", read_f=read_jsonl)\n",
    "print(f\"{len(raw_metas):,} tracks with {sum([m['duration_s'] for m in raw_metas])/60/60:,.1f}h total\")\n",
    "df_raw = pd.DataFrame.from_dict(raw_metas)\n",
    "print(df_raw.columns)\n",
    "df_raw.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_metas = read_from_s3(\"s3://suno-data/datasets/harvest/pond5_sfx/pond5_metas_v0.jsonl\", read_f=read_jsonl)\n",
    "print(f\"{len(raw_metas):,} tracks with {sum([m['duration_s'] for m in raw_metas])/60/60:,.1f}h total\")\n",
    "df_raw = pd.DataFrame.from_dict(raw_metas)\n",
    "print(df_raw.columns)\n",
    "df_raw.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k,v in df_raw.sample(n=1).iloc[0].items():\n",
    "    print(f\"{k}:{v}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}