{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, write_jsonl, read_json, write_json\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "ws = read_jsonl(\"/home/sara/task_data/whosampled_processed.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "ws[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ws2 = read_jsonl(\"/home/sara/task_data/final_whosampled_samples_11_12.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ws2[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(ws2)\n",
    "print(len(df))\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter = df[df[\"duration\"].between(60, 360, inclusive='both')]\n",
    "print(len(df_filter))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sorted = df_filter.sort_values('source_vote', ascending=False)\n",
    "df_sorted.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(df_sorted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_ids = set()\n",
    "source_to_output = {}\n",
    "for idx, row in df_sorted.iterrows():\n",
    "    output_id = row['output_id']\n",
    "    source_id = row['source_id']\n",
    "    if source_id not in source_to_output:\n",
    "        source_to_output[source_id] = []\n",
    "    if len(source_to_output[source_id]) >= 5:\n",
    "        continue\n",
    "    source_to_output[source_id].append(output_id)\n",
    "    unique_ids.add(output_id)\n",
    "    unique_ids.add(source_id)\n",
    "    if len(unique_ids) >= 50_000:\n",
    "        print(idx)\n",
    "        break\n",
    "print(len(unique_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "formated_for_scrape = []\n",
    "for id in unique_ids:\n",
    "    meta = {\"ytm_song\": {\"videoId\": id}, \"release_group_tags\": [\"whosampled\"]}\n",
    "    formated_for_scrape.append(meta)\n",
    "print(formated_for_scrape[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(formated_for_scrape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(formated_for_scrape, \"whosampled_sample_to_scrape_top_50k.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "mashup = read_jsonl(\"mashup_meta_v0.jsonl\")\n",
    "mashup_df = pd.DataFrame(mashup)\n",
    "print(len(mashup))\n",
    "print(mashup[0].keys())\n",
    "mashup_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "mashup_filter = mashup_df[\n",
    "    (mashup_df[\"mashup_strength\"].str.contains(\"strong\", na=False)) & \n",
    "    (mashup_df[\"source_ids\"].apply(len) == 2) &\n",
    "    (mashup_df[\"data_source\"] != \"whosampled_sample\")\n",
    "    #(mashup_df[\"parsing_confidence\"].fillna(0) > 0.7) &\n",
    "    #(mashup_df[\"output_match_score\"].fillna(0) > 80.0)\n",
    "]\n",
    "print(len(mashup_filter))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_ids = set()\n",
    "source_to_output = {}\n",
    "for idx, row in mashup_filter.iterrows():\n",
    "    output_id = row['output_id']\n",
    "    source_ids = row['source_ids']\n",
    "    for source_id in source_ids:\n",
    "        if source_id not in source_to_output:\n",
    "            source_to_output[source_id] = []\n",
    "        if len(source_to_output[source_id]) >= 5:\n",
    "            continue\n",
    "        source_to_output[source_id].append(output_id)\n",
    "        unique_ids.add(output_id)\n",
    "        unique_ids.add(source_id)\n",
    "    if len(unique_ids) >= 50_000:\n",
    "        print(idx)\n",
    "        break\n",
    "print(len(unique_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "mashup_filter = mashup_df[\n",
    "    ~(mashup_df[\"mashup_strength\"].str.contains(\"very weak\", na=False)) & \n",
    "    (mashup_df[\"source_ids\"].apply(len) == 2) &\n",
    "    (mashup_df[\"parsing_confidence\"] > 0.5) &\n",
    "    (mashup_df[\"data_source\"] == \"discogs\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(mashup_filter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "mashup_filter.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "mashup_filter_sorted = mashup_filter.copy()\n",
    "# Use minimum vote to ensure both are high\n",
    "mashup_filter_sorted['min_score'] = mashup_filter_sorted['source_match_sources'].apply(lambda x: min(x) if isinstance(x, list) and len(x) > 0 else 0)\n",
    "mashup_filter_sorted = mashup_filter_sorted.sort_values('min_score', ascending=False)\n",
    "mashup_filter_sorted.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, row in mashup_filter_sorted.iterrows():\n",
    "    output_id = row['output_id']\n",
    "    source_ids = row['source_ids']\n",
    "    for source_id in source_ids:\n",
    "        if source_id not in source_to_output:\n",
    "            source_to_output[source_id] = []\n",
    "        #if len(source_to_output[source_id]) >= 5:\n",
    "        #    continue\n",
    "        source_to_output[source_id].append(output_id)\n",
    "        unique_ids.add(output_id)\n",
    "        unique_ids.add(source_id)\n",
    "    if len(unique_ids) >= 50_000:\n",
    "        print(idx)\n",
    "        break\n",
    "print(len(unique_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "mashup_filter = mashup_df[\n",
    "    ~(mashup_df[\"mashup_strength\"].str.contains(\"very weak\", na=False)) & \n",
    "    (mashup_df[\"source_ids\"].apply(len) == 2) &\n",
    "    #(mashup_df[\"parsing_confidence\"] > 0.6) &\n",
    "    (mashup_df[\"data_source\"] == \"whosampled_sample\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(mashup_filter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "mashup_filter_sorted = mashup_filter.copy()\n",
    "# Use minimum vote to ensure both are high\n",
    "mashup_filter_sorted['min_score'] = mashup_filter_sorted['source_votes'].apply(lambda x: min(x) if isinstance(x, list) and len(x) > 0 else 0)\n",
    "mashup_filter_sorted = mashup_filter_sorted.sort_values('min_score', ascending=False)\n",
    "mashup_filter_sorted.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, row in mashup_filter_sorted.iterrows():\n",
    "    output_id = row['output_id']\n",
    "    source_ids = row['source_ids']\n",
    "    for source_id in source_ids:\n",
    "        if source_id not in source_to_output:\n",
    "            source_to_output[source_id] = []\n",
    "        #if len(source_to_output[source_id]) >= 5:\n",
    "        #    continue\n",
    "        source_to_output[source_id].append(output_id)\n",
    "        unique_ids.add(output_id)\n",
    "        unique_ids.add(source_id)\n",
    "    if len(unique_ids) >= 50_000:\n",
    "        print(idx)\n",
    "        break\n",
    "print(len(unique_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "formated_for_scrape = []\n",
    "for id in unique_ids:\n",
    "    meta = {\"ytm_song\": {\"videoId\": id}, \"release_group_tags\": [\"mashup\"]}\n",
    "    formated_for_scrape.append(meta)\n",
    "print(formated_for_scrape[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(formated_for_scrape, \"mashups_to_scrape_top_50k.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {},
   "source": [
    "## Get Failed YT Downloads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Load JSONL with error handling\n",
    "data = []\n",
    "errors = []\n",
    "\n",
    "with open(\"ytm_tagged_metas_test.jsonl\", \"r\") as f:\n",
    "    for line_num, line in enumerate(f, 1):\n",
    "        try:\n",
    "            data.append(json.loads(line))\n",
    "        except json.JSONDecodeError as e:\n",
    "            errors.append((line_num, str(e)))\n",
    "\n",
    "print(f\"Successfully loaded {len(data)} rows\")\n",
    "if errors:\n",
    "    print(f\"Failed to parse {len(errors)} rows:\")\n",
    "    for line_num, error in errors[:5]:  # Show first 5 errors\n",
    "        print(f\"  Line {line_num}: {error}\")\n",
    "\n",
    "df = pd.DataFrame(data)\n",
    "print(f\"\\nDataFrame shape: {df.shape}\")\n",
    "df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "failed = df[~df['success']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_bot_check = df[df['fail_message'].str.contains('sign in to confirm', case=False, na=False)]\n",
    "print(f\"Found {len(df_bot_check)} rows with 'sign in to confirm' in fail_message\")\n",
    "print(f\"Percentage: {len(df_bot_check) / len(df) * 100:.2f}%\")\n",
    "df_bot_check.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "\n",
    "class FailureCategory(Enum):\n",
    "    SUCCESS = \"success\"\n",
    "    REMOVED = \"removed\"\n",
    "    FORBIDDEN = \"403_forbidden\"\n",
    "    UNAVAILABLE = \"unavailable\"\n",
    "    TIMED_OUT = \"timed_out\"\n",
    "    BOT_CHECK = \"bot_check\"\n",
    "    COUNTRY = \"country\"\n",
    "    PRIVATE = \"private\"\n",
    "    OTHER = \"other\"\n",
    "    UNKNOWN = \"unknown\"\n",
    "\n",
    "def categorize_failure(row):\n",
    "    \"\"\"Categorize a row based on success status and failure message.\"\"\"\n",
    "    if row['success']:\n",
    "        return FailureCategory.SUCCESS\n",
    "    \n",
    "    msg = row.get('fail_message')\n",
    "    if pd.isna(msg):\n",
    "        return FailureCategory.UNKNOWN\n",
    "    \n",
    "    msg_lower = msg.lower()\n",
    "    if 'removed' in msg_lower:\n",
    "        return FailureCategory.REMOVED\n",
    "    elif 'forbidden' in msg_lower:\n",
    "        return FailureCategory.FORBIDDEN\n",
    "    elif 'unavailable' in msg_lower:\n",
    "        return FailureCategory.UNAVAILABLE\n",
    "    elif 'timed out' in msg_lower or 'timeout' in msg_lower:\n",
    "        return FailureCategory.TIMED_OUT\n",
    "    elif 'not a bot' in msg_lower or 'sign in to confirm' in msg_lower:\n",
    "        return FailureCategory.BOT_CHECK\n",
    "    elif 'your country' in msg_lower:\n",
    "        return FailureCategory.UNAVAILABLE_IN_COUNTRY\n",
    "    elif 'private' in msg_lower:\n",
    "        return FailureCategory.PRIVATE\n",
    "    else:\n",
    "        return FailureCategory.OTHER\n",
    "\n",
    "# Apply to entire dataframe\n",
    "df['failure_category'] = df.apply(categorize_failure, axis=1)\n",
    "\n",
    "# Show statistics\n",
    "category_counts = df['failure_category'].value_counts()\n",
    "print(f\"Total rows: {len(df)}\")\n",
    "print(f\"Successful: {len(df[df['success'] == True])} ({len(df[df['success'] == True])/len(df)*100:.2f}%)\")\n",
    "print(f\"Failed: {len(df[df['success'] == False])} ({len(df[df['success'] == False])/len(df)*100:.2f}%)\")\n",
    "print(\"\\nCategory breakdown:\")\n",
    "for cat, count in category_counts.items():\n",
    "    pct = count / len(df) * 100\n",
    "    print(f\"  {cat.value:20s}: {count:4d} ({pct:5.2f}%)\")\n",
    "    \n",
    "category_counts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract videoIds for retryable failures\n",
    "retryable_categories = [FailureCategory.BOT_CHECK, FailureCategory.FORBIDDEN, FailureCategory.TIMED_OUT]\n",
    "\n",
    "df_retryable = df[df['failure_category'].isin(retryable_categories)]\n",
    "\n",
    "# Extract videoId from ytm_song column\n",
    "retry_video_ids = []\n",
    "for _, row in df_retryable.iterrows():\n",
    "    ytm_song = row.get('ytm_song')\n",
    "    if isinstance(ytm_song, dict) and 'videoId' in ytm_song:\n",
    "        retry_video_ids.append(ytm_song['videoId'])\n",
    "\n",
    "print(f\"Found {len(retry_video_ids)} videos to retry:\")\n",
    "print(f\"  BOT_CHECK: {len(df_retryable[df_retryable['failure_category'] == FailureCategory.BOT_CHECK])}\")\n",
    "print(f\"  FORBIDDEN: {len(df_retryable[df_retryable['failure_category'] == FailureCategory.FORBIDDEN])}\")\n",
    "print(f\"  TIMED_OUT: {len(df_retryable[df_retryable['failure_category'] == FailureCategory.TIMED_OUT])}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "formated_for_scrape = []\n",
    "for id in retry_video_ids:\n",
    "    meta = {\"ytm_song\": {\"videoId\": id}, \"release_group_tags\": [\"mashup\"]}\n",
    "    formated_for_scrape.append(meta)\n",
    "print(formated_for_scrape[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(formated_for_scrape, \"sample_yt_ids_to_retry.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "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
}