{
 "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",
    "import emoji"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_metas = read_from_s3(\"s3://suno-data/datasets/harvest/pond5_sfx/pond5_metas.jsonl\", read_f=read_jsonl)\n",
    "print(f\"{len(raw_metas):,} tracks with {sum([m['duration'] for m in raw_metas])/60/60:,.1f}h total\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.Series([m['duration'] for m in raw_metas]).clip(upper=250).hist(bins=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_raw = pd.DataFrame.from_dict(raw_metas)\n",
    "df_raw.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Filter pond5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter = df_raw[['genre', 'description', 'name', 'author', 'details', 'id', 'duration_s', 'tags', 'related_tags', 'taxonomy', 's3_filepath']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "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": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_taxonomy(entry):\n",
    "    # If it's not a list, convert it\n",
    "    if not isinstance(entry, list):\n",
    "        if isinstance(entry, str):\n",
    "            # Convert string to list of length 1\n",
    "            entry = [entry]\n",
    "        else:\n",
    "            # Convert anything else to None\n",
    "            entry = None\n",
    "    \n",
    "    # Filter out \"no match\" from the list\n",
    "    entry = [item for item in entry if isinstance(item, str) and \"no_match\" not in item.lower()]\n",
    "    if len(entry) == 0:\n",
    "        entry = None\n",
    "    return entry\n",
    "\n",
    "# Apply the cleaning function to the taxonomy column\n",
    "df_filter['taxonomy'] = df_filter['taxonomy'].apply(clean_taxonomy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_genres(genre):\n",
    "    if isinstance(genre, str):\n",
    "        if genre.lower() == \"time lapse\":\n",
    "            genre = None\n",
    "        elif genre.lower() == \"available for musical works\":\n",
    "            genre = None\n",
    "    elif isinstance(genre, list):\n",
    "        genre = \", \". join(genre)\n",
    "    else:\n",
    "        genre = None\n",
    "\n",
    "    return genre\n",
    "\n",
    "df_filter['genre'] = df_filter['genre'].apply(filter_genres)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_description(description):\n",
    "    if isinstance(description, list):\n",
    "        description = \", \".join(description)\n",
    "    if not isinstance(description, str):\n",
    "        return None\n",
    "    return description\n",
    "\n",
    "df_filter[\"description\"] = df_filter['description'].apply(filter_description)\n",
    "df_filter['description'] = df_filter['description'].where(\n",
    "    df_filter.groupby('description')['description'].transform('count') <= 500, \n",
    "    None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_name(name):\n",
    "    if not isinstance(name, str):\n",
    "        return None\n",
    "    if name.lower() == \"also in the odyssey essentials\":\n",
    "        return None\n",
    "    if name.lower() == \"design element glitch digital\":\n",
    "        return None\n",
    "    return emoji.replace_emoji(name, replace='')\n",
    "\n",
    "df_filter['name'] = df_filter['name'].apply(filter_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_substring_after_cleaning(name, description):\n",
    "    if pd.isna(name) or pd.isna(description) or name == '':\n",
    "        return False\n",
    "    \n",
    "    # Clean both strings\n",
    "    clean_name = re.sub(r'[^a-zA-Z0-9\\s]', '', str(name)).strip().lower()\n",
    "    clean_desc = re.sub(r'[^a-zA-Z0-9\\s]', '', str(description)).strip().lower()\n",
    "    \n",
    "    return clean_name != '' and clean_name in clean_desc\n",
    "\n",
    "# Apply the condition\n",
    "df_filter.loc[df_filter.apply(lambda row: is_substring_after_cleaning(row['name'], row['description']), axis=1), 'name'] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter['id'] = df_filter['id'].str.split(\"-\", n=1).str[0]\n",
    "df_filter = df_filter.drop_duplicates(subset=['id'], keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_stereo_info(text):\n",
    "    if pd.isna(text):\n",
    "        return None\n",
    "    \n",
    "    # Search for 'stereo' or 'mono' in the text\n",
    "    match = re.search(r'\\b(stereo|mono)\\b', str(text), re.IGNORECASE)\n",
    "    \n",
    "    if match:\n",
    "        audio_type = match.group(1).lower()\n",
    "        return audio_type == 'stereo'  # True for stereo, False for mono\n",
    "    else:\n",
    "        return None  # Neither stereo nor mono found\n",
    "\n",
    "# Apply the function to create the new column\n",
    "df_filter['stereo'] = df_filter['details'].apply(extract_stereo_info)\n",
    "df_filter = df_filter.drop('details', axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def safe_process_tags(value):\n",
    "    \"\"\"Safely convert a value to a list of strings\"\"\"\n",
    "    if value is None:\n",
    "        return []\n",
    "    \n",
    "    # Handle NaN values\n",
    "    try:\n",
    "        if pd.isna(value):\n",
    "            return []\n",
    "    except (TypeError, ValueError):\n",
    "        pass\n",
    "    \n",
    "    # Handle lists\n",
    "    if isinstance(value, list):\n",
    "        return [str(item).strip() for item in value if str(item).strip() and str(item).strip().lower() != 'nan']\n",
    "    \n",
    "    # Handle strings\n",
    "    if isinstance(value, str) and value.strip() and value.strip().lower() != 'nan':\n",
    "        return [value.strip()]\n",
    "    \n",
    "    return []\n",
    "\n",
    "# Process each column separately then combine\n",
    "df_filter['tags_processed'] = df_filter['tags'].apply(safe_process_tags)\n",
    "df_filter['related_tags_processed'] = df_filter['related_tags'].apply(safe_process_tags)\n",
    "\n",
    "# Combine the processed lists\n",
    "df_filter['combined_tags'] = df_filter['tags_processed'] + df_filter['related_tags_processed']\n",
    "\n",
    "# Clean up\n",
    "df_filter = df_filter.drop(['tags', 'related_tags', 'tags_processed', 'related_tags_processed'], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter['combined_tags'] = df_filter['combined_tags'].apply(lambda x: list(set(x)) if len(x) > 0 else None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "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']\n",
    "                          .str.lower()\n",
    "                          .str.replace(r'[^a-zA-Z0-9\\s]', '', regex=True)\n",
    "                          .str.replace(r'\\s+', ' ', regex=True)\n",
    "                          .str.strip())\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['combined_tags'] = fastest_dedupe_tags(df_filter['combined_tags'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter['s3_filepath'] = df_filter['s3_filepath'].apply(lambda x: x.replace(\"{MEDIA}\", \"sfx\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(df_filter))\n",
    "df_filter.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_values = pd.Series([tag for tags in df_filter['combined_tags'] for tag in tags]).value_counts().head(50)\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "top_values.plot(kind='barh')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "### Save Filtered Metas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_path = \"s3://suno-data/datasets/harvest/pond5_sfx/pond5_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)"
   ]
  }
 ],
 "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
}