{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52a11be1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import re\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4aaec892",
   "metadata": {},
   "outputs": [],
   "source": [
    "FOLDER = \"/app2/suno/data/sara/imslp_filter/\"\n",
    "THRESHOLD = 1.0\n",
    "OUTPUT_PATH = os.path.join(FOLDER, f\"metas_v0_labeled_dupes_thresh_{THRESHOLD}.jsonl\")\n",
    "ARTIST_KEY = \"composer\"  # or artists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd15f031",
   "metadata": {},
   "outputs": [],
   "source": [
    "artist_full = read_jsonl(os.path.join(FOLDER, \"metas_raw.jsonl\"), progress=True)\n",
    "metas_v0 = read_jsonl(os.path.join(FOLDER, \"metas_v0.jsonl\"), progress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "499608f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "PAREN_PATTERN = re.compile(r\"\\([^)]*\\)\")\n",
    "BRACKET_PATTERN = re.compile(r\"\\[[^\\]]*\\]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b5b5a92",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_song_duplicates(song_entries, min_substring_length=4):\n",
    "    potential_duplicates = []\n",
    "\n",
    "    # Helper function to clean song titles (remove text in brackets)\n",
    "    def clean_title(title):\n",
    "        # Remove text in parentheses and square brackets\n",
    "        clean = PAREN_PATTERN.sub(\"\", title)\n",
    "        clean = BRACKET_PATTERN.sub(\"\", clean)\n",
    "        return clean.strip().lower()\n",
    "\n",
    "    # Preprocess all song titles - O(n) operation\n",
    "    processed_entries = []\n",
    "    exact_match_dict = defaultdict(list)\n",
    "\n",
    "    for id, title in song_entries:\n",
    "        stripped_title = title.strip()\n",
    "        cleaned_title = clean_title(stripped_title)\n",
    "\n",
    "        # Skip titles that are too short after cleaning\n",
    "        if len(cleaned_title) < min_substring_length:\n",
    "            continue\n",
    "\n",
    "        processed_entries.append((id, stripped_title, cleaned_title))\n",
    "\n",
    "        # Add to dictionary for exact match detection\n",
    "        exact_match_dict[cleaned_title].append((id, stripped_title))\n",
    "\n",
    "    # First check for exact matches - O(n) operation\n",
    "    for cleaned_title, entries in exact_match_dict.items():\n",
    "        if len(entries) > 1:\n",
    "            # Found exact matches\n",
    "            for i in range(len(entries)):\n",
    "                for j in range(i + 1, len(entries)):\n",
    "                    potential_duplicates.append((entries[i], entries[j]))\n",
    "\n",
    "    # Sort processed entries by cleaned title length to enable early termination\n",
    "    processed_entries.sort(key=lambda x: len(x[2]))\n",
    "\n",
    "    if THRESHOLD == 1.0:  # exact matches only\n",
    "        return potential_duplicates\n",
    "\n",
    "    # Then check for substring matches\n",
    "    for i in range(len(processed_entries)):\n",
    "        id1, title1, clean1 = processed_entries[i]\n",
    "\n",
    "        # Only need to check against longer strings (shorter strings can't contain this one)\n",
    "        for j in range(i + 1, len(processed_entries)):\n",
    "            id2, title2, clean2 = processed_entries[j]\n",
    "\n",
    "            # Skip exact matches (we already handled them)\n",
    "            if clean1 == clean2:\n",
    "                continue\n",
    "\n",
    "            # Skip if the length ratio is too small (prevents false positives)\n",
    "            if len(clean1) / len(clean2) < THRESHOLD:\n",
    "                continue\n",
    "\n",
    "            # Check if one is a substring of the other\n",
    "            if clean1 in clean2 or clean2 in clean1:\n",
    "                potential_duplicates.append(((id1, title1), (id2, title2)))\n",
    "\n",
    "    return potential_duplicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a4dfa6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ids_to_title = {}\n",
    "id_to_data = {}\n",
    "for song_info in artist_full:\n",
    "    if ARTIST_KEY not in song_info or \"title\" not in song_info:\n",
    "        continue\n",
    "\n",
    "    artist_info = song_info[ARTIST_KEY]\n",
    "    if isinstance(artist_info, str):\n",
    "        artist_ids = [artist_info]\n",
    "    else:\n",
    "        artist_ids = [x[\"id\"] for x in song_info[ARTIST_KEY]]\n",
    "\n",
    "    title = song_info[\"title\"]\n",
    "    id = song_info[\"id\"]\n",
    "    id_to_data[id] = song_info\n",
    "    for a_id in artist_ids:\n",
    "        if a_id not in ids_to_title:\n",
    "            ids_to_title[a_id] = []\n",
    "        ids_to_title[a_id].append((id, title))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0051ff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for id, data in id_to_data.items():\n",
    "    data[\"duplicate\"] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628b480b",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_tracks = 0\n",
    "num_dupes = 0\n",
    "dupes = []\n",
    "for a_id, titles in tqdm(ids_to_title.items()):\n",
    "    if len(titles) > 10000:\n",
    "        # print(a_id)\n",
    "        titles = titles[:10000]\n",
    "    dupes_maybe = find_song_duplicates(titles)\n",
    "    if len(dupes_maybe) > 0:\n",
    "        num_tracks += len(titles)\n",
    "        num_dupes += len(dupes_maybe)\n",
    "        for dupe_source, dupe in dupes_maybe:\n",
    "            # print(dupe_source, dupe)\n",
    "            # print(dupe_source[1], dupe[1])\n",
    "            dupes.append((dupe_source[1], dupe[1]))\n",
    "            id_to_data[dupe[0]][\"duplicate\"] = True\n",
    "print(f\"{num_tracks} tracks, {num_dupes} dupe pairs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "546233b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cleaned_metas = []\n",
    "filtered_metas = []\n",
    "for meta in metas_v0:\n",
    "    id = meta[\"id\"]\n",
    "    if id in id_to_data:\n",
    "        is_dupe = id_to_data[id][\"duplicate\"]\n",
    "        meta[\"is_artist_duplicate\"] = is_dupe\n",
    "        if not is_dupe:\n",
    "            filtered_metas.append(meta)\n",
    "    else:\n",
    "        filtered_metas.append(meta)\n",
    "    cleaned_metas.append(meta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a1c2d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_duplicates = len(metas_v0) - len(filtered_metas)\n",
    "percent_dupe = num_duplicates / len(metas_v0) * 100\n",
    "print(f\"Found {num_duplicates} duplicate titles out of {len(metas_v0)} tracks, {percent_dupe}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84585e09",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(cleaned_metas, OUTPUT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c621ed87",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = read_jsonl(\"/app2/suno/data/sara/imslp_filter/metas_raw_labeled_dupes_thresh_1.0.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f904bcf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40bc716c",
   "metadata": {},
   "outputs": [],
   "source": [
    "key_to_dupes = {}\n",
    "record_types = set()\n",
    "for out in outputs:\n",
    "    record_types.add(out[\"recording_category\"])\n",
    "    composer = out[ARTIST_KEY] if ARTIST_KEY in out else None\n",
    "    title = out[\"title\"] if \"title\" in out else None\n",
    "    if \"is_artist_duplicate\" in out and out[\"is_artist_duplicate\"]:\n",
    "        if composer is not None and title is not None:\n",
    "            source = out[\"duplicate_source\"]\n",
    "            if source not in key_to_dupes:\n",
    "                key_to_dupes[source] = []\n",
    "            key_to_dupes[source].append(out)\n",
    "            # print(composer, title, out[\"duplicate_source\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "422e23cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(record_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c73e129",
   "metadata": {},
   "outputs": [],
   "source": [
    "dupe_data = []\n",
    "\n",
    "for source, dupes in key_to_dupes.items():\n",
    "    if len(dupes) > 1:\n",
    "        artists = set([d[ARTIST_KEY] for d in dupes])\n",
    "        song_titles = set([d[\"title\"] for d in dupes])\n",
    "        ids = set([d[\"original_id\"] for d in dupes])\n",
    "        # print(artists, song_titles, len(dupes))\n",
    "        dupe_data.append((len(dupes), ids, source, artists, song_titles))\n",
    "    else:\n",
    "        print(\"uhhh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "117f016b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dupe_data.sort()\n",
    "dupe_data.reverse()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96ac356a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dupe_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dadf09a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}