{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3, list_s3_dir\n",
    "import gc\n",
    "import numpy as np\n",
    "import re\n",
    "import json\n",
    "import os\n",
    "from pathlib import Path\n",
    "import copy\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUT_DATA_DIR = \"/home/sara/sfx/v2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "other_metas = read_jsonl(\"/app2/suno/data/diffusion/sfx/v2/metas_filter_v3_add_key_bpm.jsonl\")\n",
    "\n",
    "for k, v in other_metas[0].items():\n",
    "    print(f\"{k}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_metas = read_jsonl(\"/home/sara/sfx/extreme_stems_consolidated_analysis.jsonl\")\n",
    "\n",
    "for k, v in extreme_metas[0].items():\n",
    "    print(f\"{k}: {v}, {type(v)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def keep_basic_chars(string_list):\n",
    "    return [re.sub(r'[^a-zA-Z0-9\\s\\-_]', '', string) for string in string_list]\n",
    "\n",
    "output_metas = []\n",
    "for meta in extreme_metas:\n",
    "    output_meta = {}\n",
    "    stem_name = meta[\"stem_name\"]\n",
    "    tags = meta[\"tags\"]\n",
    "    og_id = meta[\"id\"]\n",
    "    id = f\"extreme_{og_id}_{stem_name}\"\n",
    "\n",
    "    output_meta[\"id\"] = id\n",
    "    output_meta[\"tempo\"] = meta[\"tempo\"]\n",
    "    output_meta[\"key\"] = meta[\"key\"]\n",
    "    output_meta[\"duration_s_stem\"] = meta[\"duration_s\"]\n",
    "    output_meta[\"s3_path_stem\"] = meta[\"s3_path\"]\n",
    "    output_meta[\"local_path_stem\"] = meta[\"s3_path\"]\n",
    "    output_meta['tags'] = keep_basic_chars(tags)\n",
    "    output_meta['tags'].append(stem_name)\n",
    "    output_meta['tags'].append(f\"{stem_name} loop\")\n",
    "    output_meta['tags'].append(\"loop\")\n",
    "    output_meta[\"is_stem_reliable\"] = (not meta[\"is_mostly_silent\"])\n",
    "    output_meta[\"dataset\"] = \"extreme\"\n",
    "\n",
    "    output_metas.append(output_meta)\n",
    "\n",
    "# TODO: fill in duration_s and s3_path based on loop, id updated to include loop number"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in output_metas[0].items():\n",
    "    print(f\"{k}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_all_json_files(folder_path):\n",
    "    \"\"\"\n",
    "    Load all JSON files from a folder into one dictionary.\n",
    "    \n",
    "    Args:\n",
    "        folder_path (str): Path to the folder containing JSON files\n",
    "        \n",
    "    Returns:\n",
    "        dict: Combined dictionary with all JSON data\n",
    "    \"\"\"\n",
    "    combined_data = {}\n",
    "    folder = Path(folder_path)\n",
    "    \n",
    "    # Check if folder exists\n",
    "    if not folder.exists():\n",
    "        print(f\"Error: Folder {folder_path} does not exist\")\n",
    "        return {}\n",
    "    \n",
    "    # Get all JSON files in the folder\n",
    "    json_files = list(folder.glob(\"*.json\"))\n",
    "    \n",
    "    if not json_files:\n",
    "        print(f\"No JSON files found in {folder_path}\")\n",
    "        return {}\n",
    "    \n",
    "    print(f\"Found {len(json_files)} JSON files\")\n",
    "    \n",
    "    # Load each JSON file\n",
    "    for json_file in json_files:\n",
    "        try:\n",
    "            with open(json_file, 'r', encoding='utf-8') as f:\n",
    "                data = json.load(f)\n",
    "                \n",
    "            # Use filename (without extension) as key, or merge directly\n",
    "            # Option 1: Use filename as key\n",
    "            # file_key = json_file.stem\n",
    "            # combined_data[file_key] = data\n",
    "            \n",
    "            # Option 2: Merge all data directly (assuming each file contains multiple entries)\n",
    "            if isinstance(data, dict):\n",
    "                combined_data.update(data)\n",
    "            else:\n",
    "                # If data is not a dict, use filename as key\n",
    "                combined_data[json_file.stem] = data\n",
    "            \n",
    "        except json.JSONDecodeError as e:\n",
    "            print(f\"Error parsing {json_file.name}: {e}\")\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading {json_file.name}: {e}\")\n",
    "    \n",
    "    print(f\"Total entries in combined dictionary: {len(combined_data)}\")\n",
    "    return combined_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage\n",
    "folder_path = \"/home/sara/sara/processed_results\"\n",
    "all_data = load_all_json_files(folder_path)\n",
    "\n",
    "# Optional: Save combined data to a single file\n",
    "# with open(\"combined_data.json\", \"w\") as f:\n",
    "#     json.dump(all_data, f, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "loop_metas = []\n",
    "\n",
    "for meta in output_metas:\n",
    "    s3_id = meta[\"s3_path_stem\"]\n",
    "    meta_id = meta[\"id\"]\n",
    "    if s3_id in all_data and \"error\" not in all_data[s3_id]:\n",
    "        loops = all_data[s3_id][\"loops\"]\n",
    "        for idx, (uuid,loop) in enumerate(loops.items()):\n",
    "            peak_db = loop[\"loudness_metrics\"][\"peak_db\"]\n",
    "            loop_quality = loop[\"score\"]\n",
    "            duration_s = round(loop[\"duration_sec\"], 1)\n",
    "            s3_path = loop[\"s3_path\"]\n",
    "            loop_id = f\"loop{idx}\"\n",
    "\n",
    "            meta_loop = copy.deepcopy(meta)\n",
    "            meta_loop[\"peak_db\"] = peak_db\n",
    "            meta_loop[\"loop_quality\"] = loop_quality\n",
    "            meta_loop[\"duration_s\"] = duration_s\n",
    "            meta_loop[\"s3_path\"] = s3_path\n",
    "            meta_loop[\"id\"] = f\"{meta_id}_{loop_id}\"\n",
    "\n",
    "            loop_metas.append(meta_loop)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in loop_metas[0].items():\n",
    "    print(f\"{k}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_loop_metas = []\n",
    "\n",
    "num_reliable = 0\n",
    "for loop_meta in loop_metas:\n",
    "    is_loop_reliable = loop_meta[\"peak_db\"] > -15 and loop_meta[\"duration_s\"] < 15.0 and loop_meta[\"loop_quality\"] > 0.8\n",
    "    is_stem_reliable=loop_meta[\"is_stem_reliable\"]\n",
    "\n",
    "    final_loop_meta = dict(\n",
    "        id=loop_meta[\"id\"],\n",
    "        s3_filepath=loop_meta[\"s3_path\"],\n",
    "        duration_s=loop_meta[\"duration_s\"],\n",
    "        tags=loop_meta[\"tags\"],\n",
    "        dataset=loop_meta[\"dataset\"],\n",
    "        is_reliable=(is_loop_reliable and is_stem_reliable),\n",
    "        key=loop_meta[\"key\"],\n",
    "        bpm=loop_meta[\"tempo\"],\n",
    "    )\n",
    "\n",
    "    filtered_loop_metas.append(final_loop_meta)\n",
    "\n",
    "    if is_loop_reliable and is_stem_reliable:\n",
    "        num_reliable += 1\n",
    "\n",
    "print(num_reliable)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in filtered_loop_metas[0].items():\n",
    "    print(f\"{k}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(filtered_loop_metas, \"extreme_metas_v0_partial.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined = filtered_loop_metas + other_metas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "random.shuffle(combined)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(combined, \"combined_v3_w_extreme_metas_v0_partial.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "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
}