{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, write_jsonl, read_json\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "oai_mashup_data = read_jsonl(\"test_mashup_source_out_1114_all.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "oai_mashup_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract parsing confidence values (excluding Nones)\n",
    "parsing_confidences = [item['parsing_confidence'] for item in oai_mashup_data if item.get('parsing_confidence') is not None]\n",
    "\n",
    "# Extract number of source_artists (excluding Nones)\n",
    "source_artists_counts = [len(item['source_artists']) for item in oai_mashup_data if item.get('source_artists') is not None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot parsing confidence distribution\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.hist(parsing_confidences, bins=20, edgecolor='black', alpha=0.7)\n",
    "plt.xlabel('Parsing Confidence')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Distribution of Parsing Confidence')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot distribution of number of source_artists\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.hist(source_artists_counts, bins=range(min(source_artists_counts), max(source_artists_counts) + 2), edgecolor='black', alpha=0.7)\n",
    "plt.xlabel('Number of Source Artists')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Distribution of Number of Source Artists')\n",
    "plt.xticks(range(min(source_artists_counts), max(source_artists_counts) + 1))\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "keep_rows = []\n",
    "\n",
    "for o in oai_mashup_data:\n",
    "    parsing_confidence = o['parsing_confidence']\n",
    "    yt_duration = o['yt_duration']\n",
    "    if parsing_confidence < 0.2:\n",
    "        continue\n",
    "\n",
    "    if yt_duration and len(yt_duration) > 5:\n",
    "        continue\n",
    "    artists = o['source_artists']\n",
    "    songs = o['source_titles']\n",
    "\n",
    "    # normalize for comparison\n",
    "    artists_lower = [x.lower() if isinstance(x, str) else \"unknown\" for x in artists]\n",
    "    songs_lower = [x.lower() if isinstance(x, str) else \"unknown\" for x in songs]\n",
    "\n",
    "    # drop entries where both artist and song are unknown at the same index\n",
    "    filtered_artists = []\n",
    "    filtered_songs = []\n",
    "    for a, a_lower, s, s_lower in zip(artists, artists_lower, songs, songs_lower):\n",
    "        if a_lower == \"unknown\" and s_lower == \"unknown\":\n",
    "            continue\n",
    "        filtered_artists.append(a)\n",
    "        filtered_songs.append(s)\n",
    "\n",
    "    artists = filtered_artists\n",
    "    songs = filtered_songs\n",
    "\n",
    "    if len(artists) != len(songs) or len(artists) == 0 or len(artists) < 2 or len(artists) > 4:\n",
    "        continue\n",
    "\n",
    "    # update the row with cleaned lists\n",
    "    o['source_artists'] = artists\n",
    "    o['source_titles'] = songs\n",
    "\n",
    "    keep_rows.append(o)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "song_name_to_data = {}\n",
    "for row in keep_rows:\n",
    "    name = row['song_name']\n",
    "    artist = row.get(\"artists\", [\"Unknown Artist\"])\n",
    "    artist = artist[0] if artist else \"Unknown Artist\"\n",
    "    if name not in song_name_to_data:\n",
    "        song_name_to_data[name] = {}\n",
    "    if artist not in song_name_to_data[name]:\n",
    "        song_name_to_data[name][artist] = []\n",
    "    song_name_to_data[name][artist].append(row)\n",
    "\n",
    "len(song_name_to_data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mashups_1113 = read_jsonl(\"/home/sara/task_data/all_consolidated_mashups_11_13_missing_some_yt.jsonl\")\n",
    "print(len(mashups_1113))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "mashups_1113[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _count_unknowns(row):\n",
    "    artists = row.get('source_artists') or []\n",
    "    titles = row.get('source_titles') or []\n",
    "\n",
    "    def is_unknown(x):\n",
    "        return isinstance(x, str) and x.lower() == 'unknown'\n",
    "\n",
    "    return sum(1 for x in artists if is_unknown(x)) + sum(1 for x in titles if is_unknown(x))\n",
    "\n",
    "\n",
    "def _merge_rows_keep_non_null(rows, base_index):\n",
    "    \"\"\"Merge dicts in rows into a new dict, preferring values from rows[base_index].\n",
    "\n",
    "    For all other rows, only fill in keys where base has a null-ish value (None, \"\", [], {}).\n",
    "    \"\"\"\n",
    "    base = dict(rows[base_index])\n",
    "    for i, r in enumerate(rows):\n",
    "        if i == base_index:\n",
    "            continue\n",
    "        for k, v in r.items():\n",
    "            if v is None:\n",
    "                continue\n",
    "            if base.get(k) in (None, \"\", [], {}):\n",
    "                base[k] = v\n",
    "    return base\n",
    "\n",
    "\n",
    "for mu in mashups_1113:\n",
    "    oid = mu['output_id']\n",
    "    song_name = mu['song_name']\n",
    "    artist = mu.get(\"artists\", [\"Unknown Artist\"])\n",
    "    artist = artist[0] if artist else \"Unknown Artist\"\n",
    "    if oid is None:\n",
    "        if song_name not in song_name_to_data:\n",
    "            continue\n",
    "        if artist not in song_name_to_data[song_name]:\n",
    "            continue\n",
    "        potential_matches = song_name_to_data[song_name][artist]\n",
    "        if not potential_matches:\n",
    "            continue\n",
    "\n",
    "        if len(potential_matches) == 1:\n",
    "            chosen_row = dict(potential_matches[0])\n",
    "        else:\n",
    "            # choose the row with the fewest 'unknown' items across source_artists and source_titles\n",
    "            unknown_counts = [_count_unknowns(r) for r in potential_matches]\n",
    "            min_unknown = min(unknown_counts)\n",
    "            candidate_indices = [i for i, c in enumerate(unknown_counts) if c == min_unknown]\n",
    "\n",
    "            # tie-breaker: highest parsing_confidence\n",
    "            best_idx = candidate_indices[0]\n",
    "            best_conf = potential_matches[best_idx].get('parsing_confidence', 0.0)\n",
    "            for i in candidate_indices[1:]:\n",
    "                conf = potential_matches[i].get('parsing_confidence', 0.0)\n",
    "                if conf > best_conf:\n",
    "                    best_conf = conf\n",
    "                    best_idx = i\n",
    "\n",
    "            chosen_row = _merge_rows_keep_non_null(potential_matches, best_idx)\n",
    "\n",
    "        # merge chosen_row into mu: add new keys and overwrite only None values in mu\n",
    "        for k, v in chosen_row.items():\n",
    "            if k not in mu:\n",
    "                mu[k] = v\n",
    "            elif mu[k] is None and v is not None:\n",
    "                mu[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(mashups_1113)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "missing = 0\n",
    "keeping = []\n",
    "\n",
    "# fields that belong in the parsing_metadata blob\n",
    "_parsing_keys = [\n",
    "    'delimiter_found', 'mashup_keyword', 'split_parts', 'metadata',\n",
    "    'yt_video_title', 'yt_channel_title', 'yt_match_score', 'yt_duration',\n",
    "    'yt_url', 'yt_views', 'source_artists', 'source_titles',\n",
    "    'parsing_confidence', 'parsing_notes',\n",
    "]\n",
    "\n",
    "for mu in mashups_1113:\n",
    "    sid = mu['source_ids']\n",
    "    mu['parsing_metadata'] = None\n",
    "    if sid is None:\n",
    "        # still use presence of top-level source_artists to decide keeping\n",
    "        if 'source_artists' not in mu:\n",
    "            continue\n",
    "\n",
    "        # move parsing-related fields into a nested parsing_metadata dict\n",
    "        parsing_metadata = {}\n",
    "        for k in _parsing_keys:\n",
    "            if k in mu:\n",
    "                parsing_metadata[k] = mu[k]\n",
    "                del mu[k]\n",
    "        if parsing_metadata:\n",
    "            mu['parsing_metadata'] = parsing_metadata\n",
    "\n",
    "    keeping.append(mu)\n",
    "\n",
    "len(keeping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "keeping[80000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = read_jsonl(\"/home/sara/task_data/mashups_1114_sources_parsed.jsonl\")\n",
    "print(test[0].keys())\n",
    "test[80000]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "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
}