{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "## Make Consolidated DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_stems = read_jsonl(\"/home/vibert/shared/stems_captions_c686c20e.jsonl\")\n",
    "extreme_metas = read_jsonl(\"/home/sara/sfx/extreme_music_metas.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_labels(row_list):\n",
    "    output = []\n",
    "    for val in row_list:\n",
    "        name = val['label'].lower()\n",
    "        output.append(name)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "parsed_data = []\n",
    "\n",
    "for meta in extreme_metas:\n",
    "    description = meta[\"description\"]\n",
    "    genres = parse_labels(meta[\"genre\"])\n",
    "    subgenres = parse_labels(meta[\"subgenre\"])\n",
    "    instruments = parse_labels(meta[\"instruments\"])\n",
    "    moods = parse_labels(meta[\"moods\"])\n",
    "    eras = parse_labels(meta[\"eras\"])\n",
    "    keywords = parse_labels(meta[\"keywords\"])\n",
    "    description = description.split(\", \") if description is not None else None\n",
    "    bpm = meta[\"bpm\"]\n",
    "    key = meta[\"music_key\"].lower() if meta[\"music_key\"] is not None else None\n",
    "    stems_exist = meta[\"stems_avail\"]\n",
    "    stems = meta[\"stems\"]\n",
    "    parsed_row = dict(\n",
    "        id=int(meta[\"id\"]),\n",
    "        description=description,\n",
    "        genres=genres,\n",
    "        subgenres=subgenres,\n",
    "        instruments=instruments,\n",
    "        moods=moods,\n",
    "        eras=eras,\n",
    "        keywords=keywords,\n",
    "        tempo=bpm,\n",
    "        key=key,\n",
    "        stems_exist=stems_exist,\n",
    "        stem_name=stems\n",
    "    )\n",
    "    parsed_data.append(parsed_row)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(parsed_data)\n",
    "df_exploded = df.explode('stem_name')\n",
    "df_exploded['stem_name'].apply(lambda x: str(x))\n",
    "print(len(df_exploded))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_s3_path(row):\n",
    "    id = int(row[\"id\"])\n",
    "    stem_name = str(row[\"stem_name\"])\n",
    "    if id is not None and stem_name is not None:\n",
    "        s3_path = f\"s3://suno-data/datasets/harvest/extreme_music/audio/{id}/{stem_name}.mp3\"\n",
    "    else:\n",
    "        s3_path = None\n",
    "    return s3_path\n",
    "\n",
    "df_exploded['s3_path'] = df_exploded.apply(format_s3_path, axis=1)\n",
    "df_exploded.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "stems_data = []\n",
    "\n",
    "for track in extreme_stems:\n",
    "    track_id = track[\"id\"][8:]\n",
    "    stems = track.get(\"stems\", [])\n",
    "    for stem in stems:\n",
    "        local_path = stems[stem]\n",
    "        captions = track[\"stems_captions\"][stem]\n",
    "\n",
    "        caption_role = None\n",
    "        caption_keywords = None\n",
    "        caption_sonic = None\n",
    "        caption_stem = None\n",
    "        caption_description = None\n",
    "        if len(captions) > 0:\n",
    "            assert len(captions) == 5\n",
    "            caption_role = captions[0][\"caption\"]\n",
    "            caption_keywords = captions[1][\"caption\"]\n",
    "            caption_sonic = captions[2][\"caption\"]\n",
    "            caption_stem = captions[3][\"caption\"]\n",
    "            caption_description = captions[4][\"caption\"]\n",
    "        \n",
    "        data_row = dict(\n",
    "            id=int(track_id),\n",
    "            stem_name=str(stem),\n",
    "            local_path=local_path,\n",
    "            caption_role=caption_role,\n",
    "            caption_keywords=caption_keywords,\n",
    "            caption_sonic=caption_sonic,\n",
    "            caption_stem=caption_stem,\n",
    "            caption_description=caption_description,\n",
    "            duration_s=track[\"duration_s\"]\n",
    "        )\n",
    "        stems_data.append(data_row)\n",
    "\n",
    "stems_df = pd.DataFrame(stems_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = pd.merge(df_exploded, stems_df, on=['id', 'stem_name'], how='inner')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = result[result[\"stems_exist\"]]\n",
    "result = result[result['key'].notna()]\n",
    "result = result[result['tempo'].notna()]\n",
    "print(len(result))\n",
    "result = result.drop('stems_exist', axis=1)\n",
    "result.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "result.to_json('/home/sara/sfx/extreme_stems_consolidated_w_s3.jsonl', orient='records', lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Working from the Consolidated Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, read_jsonl\n",
    "import pandas as pd\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = read_jsonl('/home/sara/sfx/extreme_stems_consolidated_w_s3.jsonl')\n",
    "data_df = pd.DataFrame(data)\n",
    "data_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(data_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_by_stem_name(stem_name):\n",
    "    if \"Stem\" in str(stem_name):\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "df_test = data_df[data_df[\"stem_name\"].apply(filter_by_stem_name)]\n",
    "print(len(df_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def strip_to_alphanumeric(text):\n",
    "    stripped = re.sub(r'[^a-zA-Z0-9\\s]', '', text).lower()\n",
    "    return stripped.replace(\"gtr\", \"guitar\")\n",
    "\n",
    "df_test['stem_name'] = df_test['stem_name'].apply(strip_to_alphanumeric)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test = df_test[df_test['stem_name'].str.len() < 20]\n",
    "print(len(df_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_unique = df_test.drop_duplicates(subset=['id', 'stem_name'], keep=False)\n",
    "\n",
    "print(len(df_unique))\n",
    "df_unique = df_unique[df_unique['tempo'] < 190]\n",
    "df_unique = df_unique[df_unique['tempo'] > 40]\n",
    "print(len(df_unique))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_unique['caption_role'] = df_unique['caption_role'].apply(lambda x: x.split(\", \") if x is not None else None)\n",
    "df_unique['caption_keywords'] = df_unique['caption_keywords'].apply(lambda x: x.split(\", \") if x is not None else None)\n",
    "df_unique['caption_sonic'] = df_unique['caption_sonic'].apply(lambda x: x.split(\", \") if x is not None else None)\n",
    "df_unique['caption_stem'] = df_unique['caption_stem'].apply(lambda x: x.split(\", \") if x is not None else None)\n",
    "df_unique.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_unique.to_json('/home/sara/sfx/extreme_stems_full_tempo_key_unique_w_s3.jsonl', orient='records', lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "## Filter for Silence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import librosa\n",
    "import numpy as np\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from tqdm import tqdm\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyze_audio_volume(filepath, rms_threshold=0.01, silence_ratio_threshold=0.8):\n",
    "    \"\"\"\n",
    "    Analyze if an audio file is mostly silent or very quiet.\n",
    "    \n",
    "    Args:\n",
    "        filepath: Path to audio file\n",
    "        rms_threshold: RMS threshold below which audio is considered \"quiet\"\n",
    "        silence_ratio_threshold: Fraction of file that must be quiet to be considered \"mostly silent\"\n",
    "    \n",
    "    Returns:\n",
    "        dict: Analysis results with is_mostly_silent boolean and stats\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Load audio file\n",
    "        y, sr = librosa.load(filepath, sr=None)\n",
    "        \n",
    "        # Calculate RMS energy in overlapping frames\n",
    "        rms = librosa.feature.rms(y=y, frame_length=2048, hop_length=512)[0]\n",
    "        \n",
    "        # Calculate percentage of frames below threshold\n",
    "        quiet_frames = np.sum(rms < rms_threshold)\n",
    "        total_frames = len(rms)\n",
    "        silence_ratio = quiet_frames / total_frames if total_frames > 0 else 1.0\n",
    "        \n",
    "        # Overall volume statistics\n",
    "        max_rms = np.max(rms) if len(rms) > 0 else 0\n",
    "        mean_rms = np.mean(rms) if len(rms) > 0 else 0\n",
    "        \n",
    "        is_mostly_silent = silence_ratio >= silence_ratio_threshold\n",
    "        \n",
    "        return {\n",
    "            'filepath': filepath,\n",
    "            'is_mostly_silent': bool(is_mostly_silent),\n",
    "            'silence_ratio': float(silence_ratio),\n",
    "            'max_rms': float(max_rms),\n",
    "            'mean_rms': float(mean_rms),\n",
    "            'duration': float(len(y) / sr),\n",
    "            'status': 'success'\n",
    "        }\n",
    "        \n",
    "    except Exception as e:\n",
    "        return {\n",
    "            'filepath': filepath,\n",
    "            'is_mostly_silent': True,  # Assume problematic files are \"bad\"\n",
    "            'error': str(e),\n",
    "            'status': 'error'\n",
    "        }\n",
    "\n",
    "def filter_audio_files(filepaths, max_workers=8, rms_threshold=0.01, silence_ratio_threshold=0.8):\n",
    "    \"\"\"\n",
    "    Filter out mostly silent audio files from a list of filepaths.\n",
    "    \n",
    "    Args:\n",
    "        filepaths: List of audio file paths\n",
    "        max_workers: Number of parallel workers\n",
    "        rms_threshold: RMS threshold for quiet detection\n",
    "        silence_ratio_threshold: Minimum ratio of quiet frames to be considered mostly silent\n",
    "    \n",
    "    Returns:\n",
    "        tuple: (good_files, silent_files, analysis_results)\n",
    "    \"\"\"\n",
    "    results = []\n",
    "    good_files = []\n",
    "    silent_files = []\n",
    "    \n",
    "    print(f\"Analyzing {len(filepaths)} audio files...\")\n",
    "    \n",
    "    with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
    "        # Submit all tasks\n",
    "        future_to_filepath = {\n",
    "            executor.submit(analyze_audio_volume, fp, rms_threshold, silence_ratio_threshold): fp \n",
    "            for fp in filepaths\n",
    "        }\n",
    "        \n",
    "        # Process completed tasks with progress bar\n",
    "        for future in tqdm(as_completed(future_to_filepath), total=len(filepaths)):\n",
    "            result = future.result()\n",
    "            results.append(result)\n",
    "            \n",
    "            if result['status'] == 'success' and not result['is_mostly_silent']:\n",
    "                good_files.append(result['filepath'])\n",
    "            else:\n",
    "                silent_files.append(result['filepath'])\n",
    "    \n",
    "    return good_files, silent_files, results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "filepaths = list(df_unique['local_path'])\n",
    "print(len(filepaths))\n",
    "filepaths = filepaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "good_files, silent_files, analysis_results = filter_audio_files(\n",
    "    filepaths,\n",
    "    max_workers=100,          # Adjust based on your system\n",
    "    rms_threshold=0.01,     # Adjust sensitivity\n",
    "    silence_ratio_threshold=0.8  # 80% of file must be quiet\n",
    ")\n",
    "\n",
    "print(\"\\nResults:\")\n",
    "print(f\"Total files analyzed: {len(filepaths)}\")\n",
    "print(f\"Good files (not mostly silent): {len(good_files)}\")\n",
    "print(f\"Silent/quiet files: {len(silent_files)}\")\n",
    "\n",
    "with open('audio_analysis.json', 'w') as f:\n",
    "    json.dump(analysis_results, f, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {},
   "source": [
    "## Use Audio Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_json, read_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_analysis = read_json(\"/app2/suno/data/sara/sfx_analysis/extreme_music_loops/audio_analysis.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_metas = read_jsonl('/home/sara/sfx/extreme_stems_full_tempo_key_unique_w_s3.jsonl')\n",
    "print(len(unique_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in unique_metas:\n",
    "    m[\"is_mostly_silent\"] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "local_path_to_unique_metas = {m[\"local_path\"]:m for m in unique_metas}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "for audio_data in audio_analysis:\n",
    "    fp = audio_data[\"filepath\"]\n",
    "    is_mostly_silent = audio_data[\"is_mostly_silent\"]\n",
    "    if fp in local_path_to_unique_metas:\n",
    "        local_path_to_unique_metas[fp][\"is_mostly_silent\"] = is_mostly_silent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique = pd.DataFrame(unique_metas)\n",
    "unique.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique.to_json('/home/sara/sfx/extreme_stems_full_analysis.jsonl', orient='records', lines=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_consolidated = unique.copy()\n",
    "df_consolidated['tags'] = df_consolidated[['genres', 'subgenres', 'moods', 'eras', 'caption_role', 'caption_keywords', 'caption_sonic']].apply(\n",
    "    lambda row: list(set().union(*[lst for lst in row.dropna() if isinstance(lst, list)])) if any(isinstance(lst, list) for lst in row.dropna()) else [], axis=1\n",
    ")\n",
    "df_consolidated = df_consolidated.drop(columns=['genres', 'subgenres', 'moods', 'eras', 'caption_role', 'caption_keywords', 'caption_sonic', \"description\", \"keywords\", \"instruments\", \"caption_stem\", \"caption_description\"])\n",
    "df_consolidated.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_consolidated['is_mostly_silent'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_consolidated.to_json('/home/sara/sfx/extreme_stems_consolidated_analysis.jsonl', orient='records', lines=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_consolidated.iloc[300][\"s3_path\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_consolidated.iloc[300][\"is_mostly_silent\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}