{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "import pandas as pd\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def duration_to_seconds(duration: str) -> int:\n",
    "    \"\"\"Convert a time string like 'M:SS' or 'MM:SS' to total seconds as int.\n",
    "\n",
    "    Examples\n",
    "    --------\n",
    "    '3:45' -> 225\n",
    "    '0:09' -> 9\n",
    "    \"\"\"\n",
    "    if not duration:\n",
    "        return 0\n",
    "\n",
    "    parts = duration.strip().split(\":\")\n",
    "    if len(parts) != 2:\n",
    "        raise ValueError(f\"Invalid duration format (expected 'M:SS'): {duration!r}\")\n",
    "\n",
    "    minutes_str, seconds_str = parts\n",
    "    return int(minutes_str) * 60 + int(seconds_str)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_traxsource = read_jsonl(\"/home/sara/task_data/traxsource_metadata_raw.jsonl\")\n",
    "raw_traxsource[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "release_to_data = {}\n",
    "for entry in raw_traxsource:\n",
    "    release_id = entry.get('release_name', None)\n",
    "    release_date = entry.get('released_date', None)\n",
    "    if release_id is not None and release_date is not None:\n",
    "        release_id = f\"{release_id}_{release_date}\"\n",
    "        if release_id not in release_to_data:\n",
    "                release_to_data[release_id] = []\n",
    "        release_to_data[release_id].append(entry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(release_to_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "remix_releases = []\n",
    "\n",
    "for release_id, data in release_to_data.items():\n",
    "    if len(data) > 1:\n",
    "        songs = [d['song_name'].lower() for d in data]\n",
    "        versions = [d['version'] for d in data]\n",
    "        songs_set = list(set(songs))\n",
    "        if len(songs_set) == 1:\n",
    "            #print(release_id, songs_set, versions)\n",
    "            remix_releases.append(release_id)\n",
    "len(remix_releases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "formatted_traxsource_remix = []\n",
    "\n",
    "def is_og_song(d):\n",
    "    return len(d['version']) == 0 or \"original\" in d['version'].lower() or \"radio edit\" in d['version'].lower()\n",
    "\n",
    "for r in remix_releases:\n",
    "    data = release_to_data[r]\n",
    "    source_artists = None\n",
    "    source_title = None\n",
    "    for d in data:\n",
    "        if is_og_song(d):\n",
    "            source_artists = d['artist_names']\n",
    "            source_title = d['song_name']\n",
    "            break\n",
    "\n",
    "    for d in data:\n",
    "        if not is_og_song(d):\n",
    "            title = d['song_name'] + \" (\" + d['version'] + \")\"\n",
    "        else:\n",
    "            title = d['song_name']\n",
    "        artists = d['artist_names']\n",
    "        album_name = d['release_name']\n",
    "        label = d['label']\n",
    "        release_date = str(d['released_date'])[:4]\n",
    "        genre = d['genre']\n",
    "        bpm = d['bpm']\n",
    "        key = d['key']\n",
    "        duration = duration_to_seconds(d['length'])\n",
    "        metadata = {\n",
    "            \"title\": title,\n",
    "            \"artists\": artists,\n",
    "            \"album_name\": album_name,\n",
    "            \"source_artists\": source_artists,\n",
    "            \"source_title\": source_title,\n",
    "            \"label\": label,\n",
    "            \"release_date\": release_date,\n",
    "            \"duration\": duration,\n",
    "            \"genre\": genre,\n",
    "            \"bpm\": bpm,\n",
    "            \"key\": key,\n",
    "        }\n",
    "        formatted_traxsource_remix.append(metadata)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(formatted_traxsource_remix, \"/home/sara/task_data/traxsource_remix_metadata.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "traxsource_remix_metadata = read_jsonl(\"/home/sara/task_data/traxsource_remix_metadata.jsonl\")\n",
    "print(traxsource_remix_metadata[0].keys())\n",
    "traxsource_remix_metadata[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(traxsource_remix_metadata)  # expects a 'title' column\n",
    "\n",
    "# Filter out rows where the title contains \"radio edit\" (case-insensitive)\n",
    "df = df[~df[\"title\"].str.contains(\"radio edit\", case=False, na=False)]\n",
    "df = df[~df[\"title\"].str.contains(\"extended edit\", case=False, na=False)]\n",
    "\n",
    "# --- remix classifier ---\n",
    "\n",
    "STRONG_PATTERN = re.compile(\n",
    "    r\"\\b(remix|rework|bootleg|refix|flip)\\b\",\n",
    "    re.IGNORECASE,\n",
    ")\n",
    "\n",
    "WEAK_PATTERN = re.compile(\n",
    "    r\"\\b(edit|re[-\\s]?edit|club\\s+mix|extended\\s+mix|dub|version|alt\\s+version)\\b\",\n",
    "    re.IGNORECASE,\n",
    ")\n",
    "\n",
    "# segments inside (), [], {} and trailing dash segment\n",
    "MARKER_SEGMENT_PATTERN = re.compile(r\"[\\(\\[\\{]([^()\\[\\]{}]+)[\\)\\]\\}]\")\n",
    "\n",
    "\n",
    "def _extract_marker_segments(title: str) -> str:\n",
    "    segments = []\n",
    "\n",
    "    # ( ... ), [ ... ], { ... }\n",
    "    for m in MARKER_SEGMENT_PATTERN.finditer(title):\n",
    "        segments.append(m.group(1))\n",
    "\n",
    "    # trailing part after last hyphen, e.g. \"Song - Artist Remix\"\n",
    "    if \"-\" in title:\n",
    "        segments.append(title.split(\"-\")[-1])\n",
    "\n",
    "    return \" \".join(s.strip() for s in segments if s.strip())\n",
    "\n",
    "\n",
    "def classify_title(title: str) -> tuple[str, str]:\n",
    "    \"\"\"\n",
    "    Returns (remix_label, remix_confidence), where:\n",
    "      remix_label ∈ {\"strong\", \"weak\", \"none\"}\n",
    "      remix_confidence ∈ {\"very_strong\", \"strong\", \"weak\", \"none\"}\n",
    "    \"\"\"\n",
    "    if not isinstance(title, str):\n",
    "        title = \"\"\n",
    "    t = \" \".join(title.split())\n",
    "\n",
    "    has_strong_full = bool(STRONG_PATTERN.search(t))\n",
    "    has_weak_full = bool(WEAK_PATTERN.search(t))\n",
    "\n",
    "    marker_text = _extract_marker_segments(t)\n",
    "    has_strong_marker = bool(marker_text and STRONG_PATTERN.search(marker_text))\n",
    "    has_weak_marker = bool(marker_text and WEAK_PATTERN.search(marker_text))\n",
    "\n",
    "    # Strong keywords\n",
    "    if has_strong_full:\n",
    "        if has_strong_marker:\n",
    "            return \"strong\", \"very_strong\"  # strong keyword + markers\n",
    "        return \"strong\", \"strong\"\n",
    "\n",
    "    # Weak keywords\n",
    "    if has_weak_full:\n",
    "        if has_weak_marker:\n",
    "            return \"strong\", \"strong\"       # weak keyword + markers promoted\n",
    "        return \"weak\", \"weak\"\n",
    "\n",
    "    return None, None\n",
    "\n",
    "\n",
    "# --- apply to DataFrame ---\n",
    "\n",
    "df[\"remix_label\"], df[\"remix_confidence\"] = zip(\n",
    "    *df[\"title\"].fillna(\"\").map(classify_title)\n",
    ")\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential_remix = df[df[\"remix_confidence\"].notna()].copy()\n",
    "potential_remix.drop(columns=[\"remix_label\"], inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential_remix.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential_remix = potential_remix.reset_index(drop=True)\n",
    "\n",
    "beatport_output_jsonl_path = \"/home/sara/task_data/traxsource_potential_remix_11_17.jsonl\"\n",
    "\n",
    "beatport_records = potential_remix.to_dict(orient=\"records\")\n",
    "write_jsonl(beatport_records, beatport_output_jsonl_path)\n",
    "\n",
    "beatport_output_jsonl_path, len(beatport_records)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "traxsource = read_jsonl(\"/home/sara/task_data/traxsource_potential_remix_11_17.jsonl\")\n",
    "print(traxsource[0].keys())\n",
    "beatport = read_jsonl(\"/home/sara/task_data/beatport_potential_remix_11_17.jsonl\")\n",
    "print(beatport[0].keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(beatport))\n",
    "print(len(traxsource))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "for meta in beatport:\n",
    "    meta[\"artists\"] = meta[\"artist\"]\n",
    "    del meta[\"artist\"]\n",
    "    if meta[\"subgenre\"] is not None and len(meta[\"subgenre\"]) > 0:\n",
    "        meta[\"genre\"] = meta[\"genre\"] + \", \" + meta[\"subgenre\"]\n",
    "    del meta[\"subgenre\"]\n",
    "    meta[\"album_name\"] = None\n",
    "    meta[\"release_date\"] = meta[\"release_date\"][:4]\n",
    "    meta[\"source_artists\"] = None\n",
    "    meta[\"source_title\"] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert traxsource and beatport lists to DataFrames and merge them\n",
    "\n",
    "traxsource_df = pd.DataFrame(traxsource)\n",
    "traxsource_df[\"data_source\"] = \"traxsource\"\n",
    "\n",
    "beatport_df = pd.DataFrame(beatport)\n",
    "beatport_df[\"data_source\"] = \"beatport\"\n",
    "\n",
    "merged_df = pd.concat([traxsource_df, beatport_df], ignore_index=True)\n",
    "\n",
    "merged_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(merged_df))\n",
    "merged_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter rows where source_title and source_artists are both valid\n",
    "\n",
    "import collections.abc\n",
    "\n",
    "\n",
    "def is_valid_source_row(row: pd.Series) -> bool:\n",
    "    source_title = row.get(\"source_title\")\n",
    "    source_artists = row.get(\"source_artists\")\n",
    "\n",
    "    # source_title: non-empty string\n",
    "    if not isinstance(source_title, str) or not source_title.strip():\n",
    "        return False\n",
    "\n",
    "    # source_artists: list (or tuple) of strings, at least one non-empty\n",
    "    if not isinstance(source_artists, collections.abc.Sequence) or isinstance(source_artists, (str, bytes)):\n",
    "        return False\n",
    "\n",
    "    has_non_empty_artist = any(\n",
    "        isinstance(a, str) and a.strip() for a in source_artists\n",
    "    )\n",
    "\n",
    "    return has_non_empty_artist\n",
    "\n",
    "merged_df[\"sources_parsed\"] = merged_df.apply(is_valid_source_row, axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "who_sampled = read_jsonl(\"/home/sara/task_data/whosampled_task_labels.jsonl\")\n",
    "who_sampled_df = pd.DataFrame(who_sampled)\n",
    "who_sampled_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "who_sampled_remix = who_sampled_df[who_sampled_df[\"is_valid_remix\"]]\n",
    "who_sampled_remix.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine who_sampled_remix and merged_df\n",
    "# Keep all rows from merged_df and append whosampled remix rows\n",
    "\n",
    "# 1. Keep only selected columns from who_sampled_remix\n",
    "who_sampled_remix_small = who_sampled_remix[[\n",
    "    \"output_id\",\n",
    "    \"source_ids\",\n",
    "    \"source_votes\",\n",
    "    \"data_source\",\n",
    "]].copy()\n",
    "\n",
    "# 2. For who_sampled rows, set remix_confidence and sources_parsed\n",
    "who_sampled_remix_small[\"remix_confidence\"] = \"strong\"\n",
    "who_sampled_remix_small[\"sources_parsed\"] = True\n",
    "\n",
    "# 3. Concatenate into a single DataFrame (row-wise), keeping all of merged_df\n",
    "combined_df = pd.concat(\n",
    "    [merged_df, who_sampled_remix_small],\n",
    "    ignore_index=True,\n",
    "    sort=False,\n",
    ")\n",
    "\n",
    "combined_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save combined_df to JSONL\n",
    "\n",
    "output_path = \"/home/sara/task_data/trax_beat_who_combined_1117.jsonl\"\n",
    "combined_records = combined_df.to_dict(orient=\"records\")\n",
    "write_jsonl(combined_records, output_path)\n",
    "\n",
    "output_path, len(combined_records)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix = read_jsonl(\"/home/sara/task_data/discogs_potential_remix_11_17.jsonl\")\n",
    "print(discogs_remix[0].keys())\n",
    "discogs_remix_df = pd.DataFrame(discogs_remix)\n",
    "discogs_remix_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalize discogs_remix_df.artists to list of artist names\n",
    "\n",
    "def extract_artist_names(artists):\n",
    "    if isinstance(artists, list):\n",
    "        names = [a.get(\"name\") for a in artists if isinstance(a, dict) and a.get(\"name\")]\n",
    "        return names\n",
    "    return []\n",
    "\n",
    "discogs_remix_df[\"artists\"] = discogs_remix_df[\"artists\"].apply(extract_artist_names)\n",
    "\n",
    "discogs_remix_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop unused columns from discogs_remix_df\n",
    "cols_to_drop = [\n",
    "    \"lyrics\",\n",
    "    \"playlist_ids\",\n",
    "    \"discogs\",\n",
    "    \"musicbrainz_album\",\n",
    "    \"musicbrainz_track\",\n",
    "    \"rym_genres\",\n",
    "    \"ultimate_guitar\",\n",
    "    \"hooktheory\",\n",
    "    \"bpm\",\n",
    "    \"keywords\",\n",
    "]\n",
    "\n",
    "existing_to_drop = [c for c in cols_to_drop if c in discogs_remix_df.columns]\n",
    "discogs_remix_df = discogs_remix_df.drop(columns=existing_to_drop)\n",
    "\n",
    "discogs_remix_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix_df[\"genre\"] = discogs_remix_df[\"sos_genre\"]\n",
    "discogs_remix_df[\"duration\"] = discogs_remix_df[\"duration_s\"]\n",
    "discogs_remix_df.drop(columns=[\"sos_genre\", \"duration_s\"], inplace=True)\n",
    "discogs_remix_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix_df[\"album_name\"] = discogs_remix_df[\"album\"]\n",
    "discogs_remix_df.drop(columns=[\"album\"], inplace=True)\n",
    "discogs_remix_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix_df.drop(columns=[\"views\"], inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Keep only rows whose title contains \"remix\" (case-insensitive)\n",
    "# unless remix_confidence is \"very_strong\"\n",
    "\n",
    "remix_in_title = discogs_remix_df[\"title\"].str.contains(\"remix\", case=False, na=False)\n",
    "very_strong_conf = discogs_remix_df[\"remix_confidence\"] == \"very_strong\"\n",
    "\n",
    "filter_mask = remix_in_title & very_strong_conf\n",
    "\n",
    "discogs_remix_df = discogs_remix_df[filter_mask].copy()\n",
    "\n",
    "len(discogs_remix_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix_df[\"remix_confidence\"] = \"strong\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_remix_df[\"data_source\"] = \"discogs\"\n",
    "discogs_remix_df[\"sources_parsed\"] = False\n",
    "\n",
    "# Append to combined_df\n",
    "combined_df = pd.concat(\n",
    "    [combined_df, discogs_remix_df],\n",
    "    ignore_index=True,\n",
    "    sort=False,\n",
    ")\n",
    "\n",
    "len(combined_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = \"/home/sara/task_data/trax_beat_who_discogs_combined_1117.jsonl\"\n",
    "combined_records = combined_df.to_dict(orient=\"records\")\n",
    "write_jsonl(combined_records, output_path)\n",
    "output_path, len(combined_records)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(combined_df[~combined_df['sources_parsed']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "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
}