{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import boto3\n",
    "from tqdm import tqdm\n",
    "from suno_utils.utils.s3 import read_from_s3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_s3_files(bucket_name, prefix, max_keys: int = 100000):\n",
    "    all_files = []\n",
    "    continuation_token = None\n",
    "\n",
    "    while True:\n",
    "        # Prepare the arguments for the request\n",
    "        list_kwargs = {\n",
    "            \"Bucket\": bucket_name,\n",
    "            \"Prefix\": prefix,  # List objects under this prefix, or leave blank for all objects\n",
    "        }\n",
    "\n",
    "        if continuation_token:\n",
    "            list_kwargs[\"ContinuationToken\"] = continuation_token\n",
    "\n",
    "        # Make the request to list objects\n",
    "        response = s3.list_objects_v2(**list_kwargs)\n",
    "\n",
    "        # Collect the file keys\n",
    "        all_files += [obj[\"Key\"] for obj in response.get(\"Contents\", [])]\n",
    "\n",
    "        # Check if more results are available\n",
    "        if response.get(\"IsTruncated\"):  # True if there are more results to fetch\n",
    "            continuation_token = response[\"NextContinuationToken\"]\n",
    "        else:\n",
    "            break  # No more results to fetch\n",
    "\n",
    "    return all_files\n",
    "\n",
    "def list_s3_directories(bucket_name: str, prefix: str = \"\"):\n",
    "    \"\"\"\n",
    "    List all directories (prefixes) in an S3 bucket.\n",
    "\n",
    "    Args:\n",
    "        bucket_name (str): Name of the S3 bucket\n",
    "        prefix (str): Optional prefix to filter results (like a directory path)\n",
    "\n",
    "    Returns:\n",
    "        List[str]: List of directory paths (prefixes)\n",
    "    \"\"\"\n",
    "    s3_client = boto3.client(\"s3\")\n",
    "    directories = set()\n",
    "\n",
    "    # Use paginator to handle buckets with many objects\n",
    "    paginator = s3_client.get_paginator(\"list_objects_v2\")\n",
    "    page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter=\"/\")\n",
    "\n",
    "    # Collect all prefixes (directories)\n",
    "    for page in page_iterator:\n",
    "        # Get common prefixes (directories)\n",
    "        if \"CommonPrefixes\" in page:\n",
    "            for prefix_obj in page[\"CommonPrefixes\"]:\n",
    "                directories.add(prefix_obj[\"Prefix\"])\n",
    "\n",
    "        # Also check Contents for any directory-like objects\n",
    "        if \"Contents\" in page:\n",
    "            for obj in page[\"Contents\"]:\n",
    "                key = obj[\"Key\"]\n",
    "                # If the key contains a slash, add the directory part\n",
    "                if \"/\" in key:\n",
    "                    directory = key.rsplit(\"/\", 1)[0] + \"/\"\n",
    "                    directories.add(directory)\n",
    "\n",
    "    return sorted(list(directories))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "bucket_name = \"suno-data\"\n",
    "base_dir = \"christian/data/upsample_v4_t_5_20241018\"\n",
    "output_name = \"25hz_20241031_v1/\"\n",
    "\n",
    "# s3 client\n",
    "s3 = boto3.client(\"s3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all dir paths on s3\n",
    "dir_paths = list_s3_directories(bucket_name, f\"{base_dir}/{output_name}\")\n",
    "print(\"total dirs: \", len(dir_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from concurrent.futures import ThreadPoolExecutor\n",
    "from functools import partial\n",
    "\n",
    "def get_quality_score(dir_path, bucket_name):\n",
    "    quality_filepath = f\"s3://{bucket_name}/{dir_path}quality_scores.json\"\n",
    "    try:\n",
    "        # get quality scores from json\n",
    "        score = json.loads(read_from_s3(quality_filepath))\n",
    "        return dir_path, score\n",
    "    except Exception as e:\n",
    "        return dir_path, None\n",
    "\n",
    "quality_scores = {}\n",
    "with ThreadPoolExecutor(max_workers=64) as executor:\n",
    "    get_score_partial = partial(get_quality_score, bucket_name=bucket_name)\n",
    "    results = list(tqdm(executor.map(get_score_partial, dir_paths), total=len(dir_paths)))\n",
    "    \n",
    "quality_scores = {k: v for k, v in results if v is not None}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"total quality scores: \", len(quality_scores))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"total quality scores: \", len(quality_scores))\n",
    "\n",
    "test_id = \"christian/data/upsample_100z_v1/v2/00172f81-31b5-497f-9e6f-18782b259a10-419461\"\n",
    "result = quality_scores[list(quality_scores.keys())[0]]\n",
    "\n",
    "# turn all items into dict\n",
    "quality_scores = {k: dict(v) for k, v in quality_scores.items()}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = quality_scores[list(quality_scores.keys())[0]]\n",
    "print(result[test_id])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_filepath = \"/home/christian/code/christian/notebooks/diff_dpo/upsample_v4_t_5_20241018_25hz_20241031_v1_quality_scores.json\"\n",
    "with open(output_filepath, \"w\") as f:\n",
    "    json.dump(quality_scores, f)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env2",
   "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": 2
}