{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "import re\n",
    "import pandas as pd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_whosample_metadata = read_jsonl(\"/home/sara/who_sampled_victor/who_sampled/data/all_metadata_sara.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "covers_processed = read_jsonl(\"/home/sara/who_sampled_victor/who_sampled/data/cover_processed.jsonl\")\n",
    "samples_processed = read_jsonl(\"/home/sara/who_sampled_victor/who_sampled/data/sample_processed.jsonl\")\n",
    "remix_processed = read_jsonl(\"/home/sara/who_sampled_victor/who_sampled/data/remix_processed.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "covers_df = pd.DataFrame(covers_processed)\n",
    "samples_df = pd.DataFrame(samples_processed)\n",
    "remix_df = pd.DataFrame(remix_processed)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "covers_df['source_votes'] = covers_df['source_ids'].apply(lambda x: [xx['votes'] for xx in x])\n",
    "covers_df['source_ids'] = covers_df['source_ids'].apply(lambda x: [xx['id'] for xx in x])\n",
    "del covers_df['votes']\n",
    "covers_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill NaN is_mashup values with False\n",
    "samples_df['is_mashup'] = samples_df['is_mashup'].fillna(False)\n",
    "\n",
    "# Fill NaN votes with 0 and convert to integers\n",
    "samples_df['votes'] = samples_df['votes'].fillna(0).astype(int)\n",
    "\n",
    "samples_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process source_ids and create source_votes column\n",
    "def extract_source_data(row):\n",
    "    source_ids = row['source_ids']\n",
    "    \n",
    "    # Check if source_ids contains dicts or is already processed\n",
    "    if len(source_ids) > 1:\n",
    "        # Extract votes and ids from each element\n",
    "        votes = [xx['votes'] if isinstance(xx, dict) else 0 for xx in source_ids]\n",
    "        ids = [xx['id'] if isinstance(xx, dict) else xx for xx in source_ids]\n",
    "    else:\n",
    "        # Use the votes column value as a list\n",
    "        votes = [row['votes']]\n",
    "        # Extract the single id\n",
    "        ids = [source_ids[0]['id']] if isinstance(source_ids[0], dict) else [source_ids[0]]\n",
    "    \n",
    "    return votes, ids\n",
    "\n",
    "# Apply the function and unpack results\n",
    "samples_df[['source_votes', 'source_ids']] = samples_df.apply(\n",
    "    lambda row: pd.Series(extract_source_data(row)), axis=1\n",
    ")\n",
    "\n",
    "samples_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "del samples_df['votes']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "remix_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill NaN is_mashup values with False\n",
    "remix_df['is_mashup'] = remix_df['is_mashup'].fillna(False)\n",
    "\n",
    "# Fill NaN votes with 0 and convert to integers\n",
    "remix_df['votes'] = remix_df['votes'].fillna(0).astype(int)\n",
    "\n",
    "remix_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process source_ids and create source_votes column\n",
    "def extract_source_data_remix(row):\n",
    "    source_ids = row['source_ids']\n",
    "    \n",
    "    # Check if source_ids contains dicts or is already processed\n",
    "    if len(source_ids) > 1:\n",
    "        # Extract votes and ids from each element\n",
    "        votes = [xx['votes'] if isinstance(xx, dict) else 0 for xx in source_ids]\n",
    "        ids = [xx['id'] if isinstance(xx, dict) else xx for xx in source_ids]\n",
    "    else:\n",
    "        # Use the votes column value as a list\n",
    "        votes = [row['votes']]\n",
    "        # Extract the single id\n",
    "        ids = [source_ids[0]['id']] if isinstance(source_ids[0], dict) else [source_ids[0]]\n",
    "    \n",
    "    return votes, ids\n",
    "\n",
    "# Apply the function and unpack results\n",
    "remix_df[['source_votes', 'source_ids']] = remix_df.apply(\n",
    "    lambda row: pd.Series(extract_source_data_remix(row)), axis=1\n",
    ")\n",
    "\n",
    "remix_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "del remix_df['votes']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "remix_df['data_source'] = \"whosampled_remix\"\n",
    "del remix_df['is_remix']\n",
    "remix_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "covers_df['data_source'] = \"whosampled_cover\"\n",
    "covers_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_df['data_source'] = \"whosampled_sample\"\n",
    "samples_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "covers_df['is_valid_sample'] = False\n",
    "remix_df['is_valid_sample'] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df = pd.concat([covers_df, samples_df, remix_df], ignore_index=True)\n",
    "combined_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df['is_valid_mashup'] = combined_df['is_mashup']\n",
    "del combined_df['is_mashup']\n",
    "combined_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save combined_df to JSONL\n",
    "output_path = \"/home/sara/who_sampled_victor/who_sampled/data/all_processed_sara.jsonl\"\n",
    "write_jsonl(combined_df.to_dict('records'), output_path)\n",
    "print(f\"Saved {len(combined_df)} records to {output_path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print statistics for mashups and valid samples\n",
    "total_records = len(combined_df)\n",
    "\n",
    "num_mashups = combined_df['is_valid_mashup'].sum()\n",
    "pct_mashups = (num_mashups / total_records) * 100\n",
    "\n",
    "num_valid_samples = combined_df['is_valid_sample'].sum()\n",
    "pct_valid_samples = (num_valid_samples / total_records) * 100\n",
    "\n",
    "print(f\"Total records: {total_records:,}\")\n",
    "print(f\"\\nMashups: {num_mashups:,} ({pct_mashups:.2f}%)\")\n",
    "print(f\"Valid samples: {num_valid_samples:,} ({pct_valid_samples:.2f}%)\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "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
}