{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Select Preference Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:11:04.757824Z",
     "start_time": "2024-05-26T00:11:04.555293Z"
    }
   },
   "outputs": [],
   "source": [
    "# setup tailscale if you haven't\n",
    "# https://tailscale.com/kb/1031/install-linux\n",
    "!sudo tailscale up --accept-routes=true\n",
    "\n",
    "# setup autoload\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:11:08.392310Z",
     "start_time": "2024-05-26T00:11:04.759383Z"
    }
   },
   "outputs": [],
   "source": [
    "# pip install psycopg2-binary\n",
    "# make sure sqlalchemy is >=2\n",
    "# pip install \"sqlalchemy>=2\"\n",
    "import os\n",
    "import datetime\n",
    "from collections import defaultdict, Counter\n",
    "import json\n",
    "from urllib.parse import quote\n",
    "\n",
    "import boto3\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import sqlalchemy\n",
    "import tqdm\n",
    "from botocore.exceptions import ClientError\n",
    "from suno_analytics.preference_helper import get_preference_counts\n",
    "from suno_analytics.preference_data_selection import (\n",
    "    gather_data,\n",
    "    plot_clip_distribution,\n",
    "    parse_metadata_for_basics,\n",
    "    get_concat_clip_ids,\n",
    "    validate_preference_data,\n",
    "    run_bot_detection,\n",
    "    print_out_value_counts_nicely,\n",
    "    merge_concat_clips_with_reactions,\n",
    "    plot_clip_basic_distributions,\n",
    ")\n",
    "\n",
    "\n",
    "# setup some pandas display stuff\n",
    "pd.set_option(\"display.max_rows\", 500)\n",
    "pd.set_option(\"display.max_columns\", 500)\n",
    "pd.set_option(\"display.width\", 1000)\n",
    "\n",
    "\n",
    "def get_secret():\n",
    "    secret_name = \"rds!cluster-a3b66c33-40a7-47dd-bd6e-32b1c17c9124\"\n",
    "    region_name = \"us-east-2\"\n",
    "    # Create a Secrets Manager client\n",
    "    session = boto3.session.Session()\n",
    "    client = session.client(service_name=\"secretsmanager\", region_name=region_name)\n",
    "    try:\n",
    "        get_secret_value_response = client.get_secret_value(SecretId=secret_name)\n",
    "    except ClientError as e:\n",
    "        raise e\n",
    "    secret = get_secret_value_response[\"SecretString\"]\n",
    "    return json.loads(secret)\n",
    "\n",
    "\n",
    "my_secrets = get_secret()\n",
    "\n",
    "# alternative...\n",
    "engine = sqlalchemy.create_engine(\n",
    "    \"postgresql://postgres:%s@suno-main-pgdb-prod-analytics.cnfvffydbwvc.us-east-2.rds.amazonaws.com/suno_main\"\n",
    "    % quote(my_secrets[\"password\"])\n",
    ")\n",
    "\n",
    "\n",
    "home_dir = os.path.expanduser(\"~\")\n",
    "snow_password_path = os.path.join(home_dir, \".aws\", \"snow_pw.txt\")\n",
    "if os.path.exists(snow_password_path):\n",
    "    # !pip install snowflake\n",
    "    from snowflake.core import Root\n",
    "    from snowflake.snowpark import Session\n",
    "\n",
    "    with open(snow_password_path, \"r\") as fp:\n",
    "        fp_lines = fp.readlines()\n",
    "        snow_password = fp_lines[0].strip()\n",
    "        snow_username = fp_lines[1].strip()\n",
    "\n",
    "    CONNECTION_PARAMETERS = {\n",
    "        \"account\": \"fu90569.us-east-2.aws\",\n",
    "        \"user\": snow_username,\n",
    "        \"password\": snow_password,\n",
    "        \"role\": \"ACCOUNTADMIN\",\n",
    "        \"database\": \"SUNO_PROD\",\n",
    "        \"warehouse\": \"SUNO_PROD_LARGE\",\n",
    "        \"schema\": \"PROD\",\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validate some info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:11:08.550447Z",
     "start_time": "2024-05-26T00:11:08.397196Z"
    }
   },
   "outputs": [],
   "source": [
    "# there are 4 hr time difference between eastern time and utc\n",
    "cutoff_date = \"2024-08-26 21:00:00\"  # v4-t3 out\n",
    "# cutoff_date = (datetime.datetime.now() - datetime.timedelta(hours=2)).astimezone(datetime.timezone.utc).strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "print(cutoff_date)\n",
    "\n",
    "target_model_name = \"chirp-v3p5-engine-t-3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:11:09.349857Z",
     "start_time": "2024-05-26T00:11:08.551408Z"
    }
   },
   "outputs": [],
   "source": [
    "df_all_tables = pd.read_sql_query(\n",
    "    \"SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'\",\n",
    "    engine,\n",
    ")\n",
    "# should have all the basic table names here\n",
    "assert df_all_tables[\"table_name\"].nunique() >= 61"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query the DB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gathered_data = gather_data(engine, cutoff_date) #, filter_model_name=target_model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack the information\n",
    "bots_action_df = gathered_data[\"bots_action_df\"]\n",
    "reaction_df = gathered_data[\"reaction_df\"]\n",
    "total_clip_df = gathered_data[\"total_clip_df\"]\n",
    "playlist_clip_df = gathered_data[\"playlist_clip_df\"]\n",
    "auth_user_df = gathered_data[\"auth_user_df\"]\n",
    "discord_info_df = gathered_data[\"discord_info_df\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parse out the necessary metadata early\n",
    "total_clip_df[[\"continued_parent\", \"duration\", \"source\", \"clip_type\", \"task\"]] = pd.DataFrame(\n",
    "    total_clip_df[\"metadata\"].map(parse_metadata_for_basics).tolist(), index=total_clip_df.index\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter on versions\n",
    "clip_df = total_clip_df.copy()\n",
    "total_clip_counts = clip_df.shape[0]\n",
    "print(f\"total clips: {total_clip_counts}\")\n",
    "# check the number of audio uploads\n",
    "upload_clip_df = total_clip_df[total_clip_df[\"clip_type\"] == \"upload\"].copy()\n",
    "stem_clip_df = total_clip_df[total_clip_df[\"clip_type\"] == \"stem\"].copy()\n",
    "print(\"total without model:\", (total_clip_df[\"model_name\"] == \"\").sum())\n",
    "print(\"uploads:\", upload_clip_df.shape[0])\n",
    "print(\"stems:\", stem_clip_df.shape[0])\n",
    "\n",
    "# Call the function\n",
    "plot_clip_distribution(total_clip_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Proceed with feature engineering and cleaning up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upvoted_df = reaction_df[reaction_df[\"reaction_type\"] == \"L\"].copy()\n",
    "print(f\"number of upvoates: {upvoted_df.shape[0]:,} rows\")\n",
    "upvoted_ids = upvoted_df[\"clip_id\"]\n",
    "\n",
    "flagged_df = reaction_df[reaction_df[\"flagged\"]].copy()\n",
    "print(f\"number of flagged reports: {flagged_df.shape[0]:,} rows\")\n",
    "flagged_ids = flagged_df[\"clip_id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is probably the right way to figure out the pro user group\n",
    "pro_users = set(discord_info_df[\"user_id\"].unique())\n",
    "reaction_df[\"is_pro_user\"] = reaction_df[\"user_id\"].isin(pro_users)\n",
    "clip_df[\"is_pro_user\"] = clip_df[\"user_id\"].isin(pro_users)\n",
    "\n",
    "# this is very interesting....\n",
    "# reaction check\n",
    "print(\"Reactions fraction by pro user:\")\n",
    "print_out_value_counts_nicely(reaction_df, \"is_pro_user\")\n",
    "# clip check\n",
    "print(\"Reactions fraction by pro user:\")\n",
    "print_out_value_counts_nicely(clip_df, \"is_pro_user\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find out the stem parent ids\n",
    "stem_parent_ids = set(stem_clip_df[\"metadata\"].apply(lambda x: x.get(\"stem_from_id\", \"xxx\")))\n",
    "print(\"stem parent ids:\", len(stem_parent_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:12.019778Z",
     "start_time": "2024-05-26T00:22:57.637371Z"
    }
   },
   "outputs": [],
   "source": [
    "# add clip is in playlist feature\n",
    "clip_df[\"is_in_playlist\"] = clip_df[\"id\"].isin(playlist_clip_df[\"clip_id\"].unique())\n",
    "print(\"Clips in a splaylist:\")\n",
    "print_out_value_counts_nicely(clip_df, \"is_in_playlist\")\n",
    "clip_df[\"has_stems\"] = clip_df[\"id\"].astype(str).isin(stem_parent_ids)\n",
    "print(\"Clips has stem children:\")\n",
    "print_out_value_counts_nicely(clip_df, \"has_stems\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:17.912428Z",
     "start_time": "2024-05-26T00:23:12.021726Z"
    }
   },
   "outputs": [],
   "source": [
    "# parse the metadata for histories and types\n",
    "clip_history_df = clip_df[~clip_df[\"continued_parent\"].isna()].copy()\n",
    "# these are the direct parent's ids -- not grandparents\n",
    "has_continued_children_ids = clip_history_df[\"continued_parent\"]\n",
    "print(\n",
    "    \"clips that have children:\",\n",
    "    len(has_continued_children_ids),\n",
    "    \"\\nclips that are parents:\",\n",
    "    has_continued_children_ids.nunique(),\n",
    "    \"\\n\",\n",
    "    \"Average continues from clip = \",\n",
    "    round(len(has_continued_children_ids) / len(has_continued_children_ids.unique()), 2),\n",
    ")\n",
    "# Get value counts\n",
    "print_out_value_counts_nicely(clip_df, \"source\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:22.092835Z",
     "start_time": "2024-05-26T00:23:17.914456Z"
    }
   },
   "outputs": [],
   "source": [
    "print(\"total uploads:\", (clip_df[\"model_name\"] == \"\").sum())\n",
    "print(\"clips without request id:\", (clip_df[\"request_id\"].isna()).sum())\n",
    "# the nans are concats, we want to drop them for now\n",
    "concated_clips = clip_df[clip_df[\"clip_type\"] == \"concat\"].copy()\n",
    "non_request_clips = clip_df[clip_df[\"request_id\"].isna()].copy()\n",
    "print(\"clips without request id:\", non_request_clips.shape[0], non_request_clips[\"clip_type\"].value_counts())\n",
    "# need to kick them out...\n",
    "clip_df = clip_df[~clip_df[\"request_id\"].isna()]\n",
    "print(\n",
    "    f\"Clips without request id (concat, uploads...) frac = {concated_clips.shape[0] / total_clip_counts:.5f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:22.833148Z",
     "start_time": "2024-05-26T00:23:22.094796Z"
    }
   },
   "outputs": [],
   "source": [
    "# check the model conts\n",
    "print_out_value_counts_nicely(clip_df, \"model_name\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:26.317699Z",
     "start_time": "2024-05-26T00:23:22.835074Z"
    }
   },
   "outputs": [],
   "source": [
    "print(\"pre-filter model type clip_df shape:\", clip_df.shape)\n",
    "clip_df = clip_df[(clip_df[\"model_name\"] != \"chirp-v3-5\") & (clip_df[\"model_name\"] != \"chirp-v3-0\")]\n",
    "print(\"post-filter model type clip_df shape:\", clip_df.shape)\n",
    "print_out_value_counts_nicely(clip_df, \"model_name\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "concated_clips = merge_concat_clips_with_reactions(concated_clips, reaction_df)\n",
    "# TODO: why so many clips are concats without plays??? -- oh probably they concat multiple times?\n",
    "print(\"All concats\", concated_clips.shape[0])\n",
    "concated_clips = concated_clips[concated_clips[\"reaction_play_count\"] > 0]\n",
    "print(\"total concats with plays\", concated_clips.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:38.974479Z",
     "start_time": "2024-05-26T00:23:34.385507Z"
    }
   },
   "outputs": [],
   "source": [
    "concat_clips_ids = get_concat_clip_ids(concated_clips, clip_df, upload_clip_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:40.236690Z",
     "start_time": "2024-05-26T00:23:39.995713Z"
    }
   },
   "outputs": [],
   "source": [
    "# set user number of clips generated\n",
    "clip_df[\"user_n_clips\"] = clip_df[\"user_id\"].map(clip_df[\"user_id\"].value_counts())\n",
    "print(clip_df[\"user_n_clips\"].describe())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:44.826860Z",
     "start_time": "2024-05-26T00:23:40.238330Z"
    }
   },
   "outputs": [],
   "source": [
    "# add upvoted column\n",
    "clip_df[\"upvoted\"] = clip_df[\"id\"].isin(upvoted_ids)\n",
    "print(\n",
    "    \"has upvoted\",\n",
    "    clip_df[\"upvoted\"].value_counts(),\n",
    "    clip_df[\"upvoted\"].value_counts(normalize=True),\n",
    "    (clip_df[\"upvote_count\"] >= 1).value_counts(normalize=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:48.404433Z",
     "start_time": "2024-05-26T00:23:44.857788Z"
    }
   },
   "outputs": [],
   "source": [
    "disliked_ids = reaction_df[reaction_df[\"reaction_type\"] == \"D\"][\"clip_id\"].unique()\n",
    "\n",
    "clip_df[\"downvoted\"] = clip_df[\"id\"].isin(disliked_ids)\n",
    "print(\"downvoted fraction by category:\")\n",
    "print_out_value_counts_nicely(clip_df, \"downvoted\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:23:56.590022Z",
     "start_time": "2024-05-26T00:23:48.405668Z"
    }
   },
   "outputs": [],
   "source": [
    "# add continued column -- uuid and str are not compatible X.x\n",
    "clip_df[\"has_continued\"] = clip_df[\"id\"].astype(str).isin(set(list(has_continued_children_ids)))\n",
    "print(\"has_continued fraction by category:\")\n",
    "print_out_value_counts_nicely(clip_df, \"has_continued\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:04.704306Z",
     "start_time": "2024-05-26T00:23:56.591353Z"
    }
   },
   "outputs": [],
   "source": [
    "# add concat column\n",
    "clip_df[\"part_of_concat\"] = clip_df[\"id\"].astype(str).isin(concat_clips_ids)\n",
    "print(\"part_of_concat fraction by category:\")\n",
    "print_out_value_counts_nicely(clip_df, \"part_of_concat\")\n",
    "\n",
    "print(\"\\nModel distribution for part_of_concat clips:\")\n",
    "for model, fraction in (\n",
    "    clip_df[clip_df[\"part_of_concat\"]][\"model_name\"].value_counts(normalize=True).items()\n",
    "):\n",
    "    print(f\"{model}: {fraction:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:09.954233Z",
     "start_time": "2024-05-26T00:24:04.705554Z"
    }
   },
   "outputs": [],
   "source": [
    "# verify bots action are all non-empty\n",
    "action_mask = (\n",
    "    bots_action_df[\"download_audio_count\"]\n",
    "    + bots_action_df[\"download_video_count\"]\n",
    "    + bots_action_df[\"download_audio_wav_count\"]\n",
    "    # + bots_action_df[\"share_count\"] # will remove share cause it can be negative, just can be...\n",
    ") >= 1\n",
    "has_action_ids = set(i for i in bots_action_df[action_mask][\"clip_id\"].unique())\n",
    "clip_df[\"has_action\"] = clip_df[\"id\"].isin(has_action_ids)\n",
    "print(\"has_action fraction by category:\")\n",
    "print_out_value_counts_nicely(clip_df, \"has_action\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:14.579841Z",
     "start_time": "2024-05-26T00:24:09.955568Z"
    }
   },
   "outputs": [],
   "source": [
    "# add downvoted column\n",
    "clip_df[\"flagged\"] = clip_df[\"id\"].isin(flagged_ids)\n",
    "print(\"flagged fraction by category:\")\n",
    "print_out_value_counts_nicely(clip_df, \"flagged\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_df[\"deleted\"] = clip_df[\"is_deleted\"]\n",
    "print(\"deleted fraction by category:\")\n",
    "print_out_value_counts_nicely(clip_df, \"deleted\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:15.382315Z",
     "start_time": "2024-05-26T00:24:14.581073Z"
    }
   },
   "outputs": [],
   "source": [
    "# This is probably the most important cell of this notebook -- what are good labels, and not having good label makes it a bad label\n",
    "must_be_positive_mask = (\n",
    "    (clip_df[\"upvoted\"])\n",
    "    | (clip_df[\"has_action\"])\n",
    "    | (clip_df[\"part_of_concat\"])\n",
    "    | (clip_df[\"is_in_playlist\"])\n",
    ")\n",
    "must_be_not_negative_mask = (~clip_df[\"downvoted\"]) & (~clip_df[\"deleted\"]) & (~clip_df[\"flagged\"])\n",
    "must_be_negative_mask = (clip_df[\"downvoted\"]) | (clip_df[\"flagged\"]) | (clip_df[\"deleted\"])\n",
    "total_clips_count = clip_df.shape[0]\n",
    "must_be_positive_count = sum(must_be_positive_mask)\n",
    "definitely_not_negative_count = sum(must_be_not_negative_mask)\n",
    "must_be_negative_count = sum(must_be_negative_mask)\n",
    "\n",
    "print(\n",
    "    f\"Total clips: {total_clips_count:,}\\n\"\n",
    "    f\"Must be positive: {must_be_positive_count:,} ({must_be_positive_count/total_clips_count:.2%})\\n\"\n",
    "    f\"Definitely not negative: {definitely_not_negative_count:,} ({definitely_not_negative_count/total_clips_count:.2%})\\n\"\n",
    "    f\"Must be negative: {must_be_negative_count:,} ({must_be_negative_count/total_clips_count:.2%})\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:31.095990Z",
     "start_time": "2024-05-26T00:24:15.383572Z"
    }
   },
   "outputs": [],
   "source": [
    "mask = must_be_positive_mask & must_be_not_negative_mask\n",
    "total_unique_requests = clip_df[\"request_id\"].nunique()\n",
    "liked_requests = clip_df[mask][\"request_id\"].unique()  # requests with at least 1 like\n",
    "unliked_requests = clip_df[~mask][\"request_id\"].unique()  # requests without like\n",
    "has_liked_requests = set(liked_requests).intersection(\n",
    "    set(unliked_requests)\n",
    ")  # the request must have 1 like and one without like\n",
    "print(f\"Liked requests: {len(liked_requests):,}\")\n",
    "print(f\"Not liked requests: {len(unliked_requests):,}\")\n",
    "print(f\"Requests with preference paired generations: {len(has_liked_requests):,}\")\n",
    "print(f\"Percentage of total unique requests: {len(has_liked_requests) / total_unique_requests:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# introduce a negative preference count\n",
    "has_disliked_half_requests = clip_df[must_be_negative_mask][\n",
    "    \"request_id\"\n",
    "].unique()  # requests with at least 1 dislike\n",
    "not_have_disliked_requests = clip_df[~must_be_negative_mask][\n",
    "    \"request_id\"\n",
    "].unique()  # request without dislike\n",
    "has_disliked_requests = set(has_disliked_half_requests).intersection(\n",
    "    set(not_have_disliked_requests)\n",
    ")  # the request must have 1 dislike and one without dislike\n",
    "print(f\"Disliked requests: {len(has_disliked_half_requests):,}\")\n",
    "print(f\"Not disliked requests: {len(not_have_disliked_requests):,}\")\n",
    "print(f\"Requests with preference paired generations: {len(has_disliked_requests):,}\")\n",
    "print(f\"Percentage of total unique requests: {len(has_disliked_requests) / total_unique_requests:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:31.099372Z",
     "start_time": "2024-05-26T00:24:31.097244Z"
    }
   },
   "outputs": [],
   "source": [
    "requests = has_liked_requests.union(has_disliked_requests)\n",
    "print(f\"Total selected pairs of requests: {len(requests):,}\")\n",
    "print(f\"Percentage of total unique requests: {len(requests) / total_unique_requests:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:31.239254Z",
     "start_time": "2024-05-26T00:24:31.100389Z"
    }
   },
   "outputs": [],
   "source": [
    "# this used to be a terrible bug...X.x\n",
    "assert mask.shape[0] == clip_df.shape[0]\n",
    "clip_df[\"pos_preference\"] = mask\n",
    "clip_df[\"neg_preference\"] = must_be_negative_mask\n",
    "# note that this is along the same row, so a positive clip can't be negative\n",
    "clip_df[\"diff_preference\"] = clip_df[\"pos_preference\"].astype(int) - clip_df[\"neg_preference\"].astype(\n",
    "    int\n",
    ")\n",
    "print(\"Difference in preference counts:\")\n",
    "value_counts = clip_df[\"diff_preference\"].value_counts()\n",
    "total = value_counts.sum()\n",
    "for value, count in value_counts.items():\n",
    "    fraction = count / total\n",
    "    print(f\"{value}: {count:,} ({fraction:.2%})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:37.322031Z",
     "start_time": "2024-05-26T00:24:31.240829Z"
    }
   },
   "outputs": [],
   "source": [
    "# creation of interesting_clips\n",
    "interesting_clips = clip_df[clip_df[\"request_id\"].isin(requests)].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interesting_clips = interesting_clips.sort_values(by=[\"request_id\", \"diff_preference\"]).reset_index()\n",
    "interesting_clips[[\"request_id\", \"pos_preference\", \"neg_preference\", \"diff_preference\"]].head(n=6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is a mix now\n",
    "value_counts = interesting_clips[\"diff_preference\"].value_counts()\n",
    "total = value_counts.sum()\n",
    "for value, count in value_counts.items():\n",
    "    fraction = count / total\n",
    "    print(f\"Difference {value}: {count:,} ({fraction:.2%})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_series = interesting_clips[\"diff_preference\"].diff()\n",
    "value_counts = diff_series[1::2].value_counts()\n",
    "total = value_counts.sum()\n",
    "for value, count in value_counts.items():\n",
    "    fraction = count / total\n",
    "    print(f\"Value {value}: {count:,} ({fraction:.2%})\")\n",
    "# 1 is pos, not neg pair or nothing, neg; 2 is pos / neg (hence the larger difference)\n",
    "# there are only two values for this positive pair\n",
    "assert diff_series[1::2].nunique() == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:38.960130Z",
     "start_time": "2024-05-26T00:24:37.323369Z"
    }
   },
   "outputs": [],
   "source": [
    "# assign the labels now\n",
    "interesting_clips[\"preference\"] = interesting_clips.index % 2 == 1\n",
    "# get df of requests -- let's move on!\n",
    "print(f\"Number of unique request_ids: {interesting_clips['request_id'].nunique():,}\")\n",
    "print(f\"Number of unique ids: {interesting_clips['id'].nunique():,}\")\n",
    "validate_preference_data(interesting_clips)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:43.332222Z",
     "start_time": "2024-05-26T00:24:43.166461Z"
    }
   },
   "outputs": [],
   "source": [
    "# # listen to some pairs\n",
    "# test_requests = interesting_clips[\"request_id\"].sample(10)\n",
    "\n",
    "# for i in range(1):\n",
    "#     rows = interesting_clips[interesting_clips[\"request_id\"] == test_requests.iloc[i]]\n",
    "#     assert rows.shape[0] == 2\n",
    "#     # Audio.from_s3(f\"s3://suno-data-uploads/studio/uploads/{row['s3_id']}.mp3\").play()\n",
    "#     # sort by likes\n",
    "#     rows = rows.sort_values(\"upvoted\", ascending=True)\n",
    "#     print(rows.iloc[0][\"prompt_text\"])\n",
    "#     print(rows.iloc[0][\"metadata\"])\n",
    "#     for _, row in rows.iterrows():\n",
    "#         print(row[\"id\"], row[\"preference\"], row[\"upvoted\"])\n",
    "#         Audio.from_s3(\n",
    "#             f\"s3://suno-data-uploads/studio/uploads/{row['s3_id']}.mp3\"\n",
    "#         ).play()\n",
    "#         with open_from_s3(\n",
    "#             f\"s3://suno-data-uploads/studio/uploads/{row['s3_id']}.npz\", as_binary=True\n",
    "#         ) as f:\n",
    "#             # read numpy array\n",
    "#             npz_a = np.load(f)\n",
    "#             if \"v1_raw\" in npz_a:\n",
    "#                 a = np.load(f)[\"v1_raw\"]\n",
    "#             elif \"v3.0_raw\" in npz_a:\n",
    "#                 a = np.load(f)[\"v3.0_raw\"]\n",
    "#             else:\n",
    "#                 print(\"npz_a\", npz_a)\n",
    "#                 raise ValueError\n",
    "#             print(a.shape)\n",
    "#     print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Further cuts and selections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# need the reaction play counts\n",
    "# Filter reaction_df for relevant clip_ids\n",
    "partial_reaction_df = reaction_df[reaction_df[\"clip_id\"].isin(set(interesting_clips[\"id\"]))].copy()\n",
    "\n",
    "# Calculate total play counts\n",
    "total_play_counts = partial_reaction_df.groupby(\"clip_id\")[\"play_count\"].sum().reset_index()\n",
    "total_play_counts = total_play_counts.rename(\n",
    "    columns={\"clip_id\": \"id\", \"play_count\": \"reaction_play_count\"}\n",
    ")\n",
    "\n",
    "# Calculate pro user play counts\n",
    "pro_play_counts = (\n",
    "    partial_reaction_df[partial_reaction_df[\"is_pro_user\"]]\n",
    "    .groupby(\"clip_id\")[\"play_count\"]\n",
    "    .sum()\n",
    "    .reset_index()\n",
    ")\n",
    "pro_play_counts = pro_play_counts.rename(\n",
    "    columns={\"clip_id\": \"id\", \"play_count\": \"reaction_pro_play_count\"}\n",
    ")\n",
    "\n",
    "# Merge with user_intersting_clips\n",
    "interesting_clips = interesting_clips.merge(total_play_counts, on=\"id\", how=\"left\")\n",
    "interesting_clips = interesting_clips.merge(pro_play_counts, on=\"id\", how=\"left\")\n",
    "\n",
    "print(f\"Number of interesting clips: {len(interesting_clips):,}\")\n",
    "# Get unique counts for request_id and id\n",
    "unique_request_ids = interesting_clips[\"request_id\"].nunique()\n",
    "unique_clip_ids = interesting_clips[\"id\"].nunique()\n",
    "\n",
    "# Print the results in a formatted manner\n",
    "print(\"Unique request and clip counts in interesting_clips:\")\n",
    "print(f\"{'Request IDs:':<15} {unique_request_ids:,}\")\n",
    "print(f\"{'Clip IDs:':<15} {unique_clip_ids:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:43.737067Z",
     "start_time": "2024-05-26T00:24:43.563216Z"
    }
   },
   "outputs": [],
   "source": [
    "preference_counts = interesting_clips.groupby(\"batch_index\")[\"preference\"].value_counts()\n",
    "total_counts = preference_counts.groupby(level=0).sum()\n",
    "\n",
    "print(\"Preference counts and fractions by batch index:\")\n",
    "print(\"-\" * 50)\n",
    "for batch_index in [0, 1]:\n",
    "    print(f\"Batch Index: {batch_index}\")\n",
    "    for preference in [False, True]:\n",
    "        count = preference_counts[batch_index, preference]\n",
    "        fraction = count / total_counts[batch_index]\n",
    "        print(f\"  Preference {preference}: Count: {count:,} Fraction: {fraction:.2%}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:44.304573Z",
     "start_time": "2024-05-26T00:24:43.973218Z"
    }
   },
   "outputs": [],
   "source": [
    "print_out_value_counts_nicely(interesting_clips, \"model_name\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:44.533983Z",
     "start_time": "2024-05-26T00:24:44.305722Z"
    }
   },
   "outputs": [],
   "source": [
    "print(\"Time Validation:\")\n",
    "print(\"-\" * 20)\n",
    "print(\"Interesting Clips:\")\n",
    "print(f\"  Earliest: {interesting_clips['created_at'].min()}\")\n",
    "print(f\"  Latest:   {interesting_clips['created_at'].max()}\")\n",
    "print(\"\\nAll Clips:\")\n",
    "print(f\"  Earliest: {clip_df['created_at'].min()}\")\n",
    "print(f\"  Latest:   {clip_df['created_at'].max()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:50.674838Z",
     "start_time": "2024-05-26T00:24:46.366677Z"
    }
   },
   "outputs": [],
   "source": [
    "# make sure we sort here before proceed\n",
    "interesting_clips = interesting_clips.sort_values(by=[\"request_id\", \"preference\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:24:51.600621Z",
     "start_time": "2024-05-26T00:24:50.676186Z"
    }
   },
   "outputs": [],
   "source": [
    "# Calculate the ratio of preferred clips to total clips for each model\n",
    "clip_df_model_name_value_counts = clip_df[\"model_name\"].value_counts()\n",
    "preference_ratio = (\n",
    "    interesting_clips[interesting_clips[\"preference\"]][\"model_name\"].value_counts()\n",
    "    / clip_df_model_name_value_counts\n",
    ")\n",
    "\n",
    "# Print the results in a formatted manner\n",
    "print(\"Ratio of preferred clips to total clips for each model:\")\n",
    "print(\"-\" * 60)\n",
    "for model, ratio in preference_ratio.items():\n",
    "    n = clip_df_model_name_value_counts[model]\n",
    "    uncertainty = (ratio * (1 - ratio) / n) ** 0.5\n",
    "    print(f\"{model:<30} {ratio:.2%} ± {uncertainty:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:00.484734Z",
     "start_time": "2024-05-26T00:25:00.377280Z"
    }
   },
   "outputs": [],
   "source": [
    "# from suno_analytics.preference_data_selection import parse_for_tag, parse_for_one_box\n",
    "# user_intersting_clips[\"tags\"] = user_intersting_clips[\"metadata\"].apply(parse_for_tag)\n",
    "# user_intersting_clips[\"is_onebox\"] = user_intersting_clips[\"metadata\"].apply(parse_for_one_box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:00.893917Z",
     "start_time": "2024-05-26T00:25:00.485775Z"
    }
   },
   "outputs": [],
   "source": [
    "user_compare_mask = (\n",
    "    (user_intersting_clips[\"created_at\"] >= cutoff_date)\n",
    "    # & (\n",
    "    #     (user_intersting_clips[\"model_name\"].str.startswith(\"chirp-v3p5-engine-t\"))\n",
    "    #     | (user_intersting_clips[\"model_name\"].str.startswith(\"chirp-v3p5-engine-s\"))\n",
    "    # )\n",
    "    # & (~user_intersting_clips[\"is_pro_user\"])\n",
    "    # & (~user_intersting_clips[\"is_onebox\"])\n",
    ")\n",
    "# # this is fucked up sometimes one box doesn't give prompt to one generation\n",
    "extra_compare_mask = user_intersting_clips[user_compare_mask][\"request_id\"].isin(\n",
    "    user_intersting_clips[user_compare_mask][\"request_id\"]\n",
    "    .value_counts()\n",
    "    .index[user_intersting_clips[user_compare_mask][\"request_id\"].value_counts() == 2]\n",
    ")\n",
    "\n",
    "user_compare_mask = user_compare_mask & extra_compare_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:01.140472Z",
     "start_time": "2024-05-26T00:25:00.895575Z"
    }
   },
   "outputs": [],
   "source": [
    "user_intersting_clips_3p5 = user_intersting_clips[user_compare_mask].reset_index().copy()\n",
    "\n",
    "\n",
    "def modify_model_name(model_name, metadata):\n",
    "    if model_name.startswith(\"chirp-v3p5-engine-t\") or model_name.startswith(\"chirp-v3p5-engine-s\"):\n",
    "        if \"param_experiment\" in metadata:\n",
    "            exp = metadata.get(\"param_experiment\", \"\")\n",
    "            if exp:\n",
    "                return f\"{model_name}_{exp}\"\n",
    "    return model_name\n",
    "\n",
    "\n",
    "user_intersting_clips_3p5[\"model_name\"] = user_intersting_clips_3p5.apply(\n",
    "    lambda row: modify_model_name(row[\"model_name\"], row[\"metadata\"]), axis=1\n",
    ")\n",
    "user_intersting_clips_3p5 = user_intersting_clips_3p5.sort_values(by=[\"request_id\", \"preference\"])\n",
    "print(user_intersting_clips_3p5.shape)\n",
    "model_counts = user_intersting_clips_3p5[\"model_name\"].value_counts()\n",
    "model_fracs = model_counts / model_counts.sum()\n",
    "\n",
    "print(\"Model Name Value Counts and Fractions:\")\n",
    "print_out_value_counts_nicely(user_intersting_clips_3p5, \"model_name\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:02.009450Z",
     "start_time": "2024-05-26T00:25:01.523515Z"
    }
   },
   "outputs": [],
   "source": [
    "get_preference_counts(user_intersting_clips_3p5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:02.332947Z",
     "start_time": "2024-05-26T00:25:02.010694Z"
    }
   },
   "outputs": [],
   "source": [
    "print(\"first gen\")\n",
    "first_gen_slice_df = user_intersting_clips_3p5[\n",
    "    (user_intersting_clips_3p5[\"continued_parent\"].isna())\n",
    "].copy()\n",
    "if first_gen_slice_df.shape[0] > 0:\n",
    "    get_preference_counts(\n",
    "        user_intersting_clips_3p5[(user_intersting_clips_3p5[\"continued_parent\"].isna())],\n",
    "        title_name=\"first generation\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:02.574659Z",
     "start_time": "2024-05-26T00:25:02.334214Z"
    }
   },
   "outputs": [],
   "source": [
    "print(\"is continue\")\n",
    "get_preference_counts(\n",
    "    user_intersting_clips_3p5[(~user_intersting_clips_3p5[\"continued_parent\"].isna())],\n",
    "    \"is continue\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clean up SHIT\n",
    "\n",
    "to get the right play conts, we need the right df..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:08.095035Z",
     "start_time": "2024-05-26T00:25:07.782738Z"
    }
   },
   "outputs": [],
   "source": [
    "def unpack_dict(x):\n",
    "    if v := concat_clips_ids.get(str(x)):\n",
    "        return v\n",
    "    else:\n",
    "        return {\n",
    "            \"total_start_s\": None,\n",
    "            \"total_clip_s\": None,\n",
    "            \"concat_play_counts\": None,\n",
    "            \"concat_in_playlist\": None,\n",
    "            \"concat_likes\": None,\n",
    "            \"concat_dislikes\": None,\n",
    "        }\n",
    "\n",
    "\n",
    "extra_cols = user_intersting_clips[\"id\"].apply(unpack_dict)\n",
    "extra_cols_df = pd.DataFrame.from_records(extra_cols.values, index=extra_cols.index)\n",
    "user_intersting_clips[\n",
    "    [\n",
    "        \"total_start_s\",\n",
    "        \"total_clip_s\",\n",
    "        \"concat_play_counts\",\n",
    "        \"concat_in_playlist\",\n",
    "        \"concat_likes\",\n",
    "        \"concat_dislikes\",\n",
    "    ]\n",
    "] = extra_cols_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:08.361849Z",
     "start_time": "2024-05-26T00:25:08.199583Z"
    }
   },
   "outputs": [],
   "source": [
    "pos_too_much_data_mask = (\n",
    "    (user_intersting_clips[\"preference\"])\n",
    "    & (\n",
    "        (user_intersting_clips[\"reaction_play_count\"] >= 2)  # single play is super catchy\n",
    "        | (user_intersting_clips[\"concat_play_counts\"] >= 2)  # or the concat play is super catchy\n",
    "    )\n",
    "    # & (user_intersting_clips[\"user_n_clips\"] >= 40)\n",
    "    # & (user_intersting_clips[\"continued_parent\"].isna())\n",
    ")\n",
    "neg_too_much_data_mask = (\n",
    "    (~user_intersting_clips[\"preference\"])\n",
    "    & (user_intersting_clips[\"reaction_play_count\"] >= 1)  # single play is super catchy\n",
    "    # & (user_intersting_clips[\"user_n_clips\"] >= 40)\n",
    "    # & (user_intersting_clips[\"continued_parent\"].isna())\n",
    ")\n",
    "# Calculate and print the proportion of data that meets our criteria\n",
    "pos_proportion = pos_too_much_data_mask.sum() / user_intersting_clips.shape[0] * 2\n",
    "print(f\"Positive proportion of data meeting criteria: {pos_proportion:.2%}\")\n",
    "neg_proportion = neg_too_much_data_mask.sum() / user_intersting_clips.shape[0] * 2\n",
    "print(f\"Negative proportion of data meeting criteria: {neg_proportion:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:08.523643Z",
     "start_time": "2024-05-26T00:25:08.367348Z"
    }
   },
   "outputs": [],
   "source": [
    "final_good_enough_requests = set(\n",
    "    user_intersting_clips[pos_too_much_data_mask][\"request_id\"].unique()\n",
    ").intersection(set(user_intersting_clips[neg_too_much_data_mask][\"request_id\"].unique()))\n",
    "final_interesting_clips = user_intersting_clips[\n",
    "    user_intersting_clips[\"request_id\"].isin(final_good_enough_requests)\n",
    "].copy()\n",
    "# Get the value counts of model_name for preferred clips\n",
    "model_counts = final_interesting_clips[final_interesting_clips[\"preference\"]][\n",
    "    \"model_name\"\n",
    "].value_counts()\n",
    "\n",
    "# Print the results in a nicely formatted way\n",
    "total_count = model_counts.sum()\n",
    "print(\"Model Name Value Counts for Preferred Clips:\")\n",
    "print(\"-\" * 70)\n",
    "print(f\"{'Model':<30} {'Count':>10} {'Fraction':>15}\")\n",
    "print(\"-\" * 70)\n",
    "for model, count in model_counts.items():\n",
    "    fraction = count / total_count\n",
    "    print(f\"{model:<30} {count:>10,d} {fraction:>15.2%}\")\n",
    "print(\"-\" * 70)\n",
    "print(f\"{'Total':<30} {total_count:>10,d} {1:>15.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:09.198504Z",
     "start_time": "2024-05-26T00:25:08.885507Z"
    }
   },
   "outputs": [],
   "source": [
    "validate_preference_data(final_interesting_clips)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the number of rows for final_interesting_clips with the specific model\n",
    "row_count = final_interesting_clips[final_interesting_clips[\"model_name\"] == target_model_name].shape[0]\n",
    "\n",
    "# Print the row count in a nicely formatted way\n",
    "print(f\"Number of rows in final_interesting_clips for model {target_model_name}:\")\n",
    "print(f\"{row_count:,}\")\n",
    "print(\"done\", final_interesting_clips.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# For faster processing once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-26T00:25:10.265241Z",
     "start_time": "2024-05-26T00:25:09.934743Z"
    }
   },
   "outputs": [],
   "source": [
    "# Calculate the number of unique users\n",
    "total_unique_users = clip_df[\"user_id\"].nunique()\n",
    "\n",
    "# Print the result in a nicely formatted way\n",
    "print(\"Total Unique Users:\")\n",
    "print(\"-\" * 20)\n",
    "print(f\"{total_unique_users:,}\")\n",
    "print(\"-\" * 20)\n",
    "\n",
    "# This can take a while cause we have a lot of users...\n",
    "# query = \"\"\"\n",
    "# SELECT *\n",
    "# FROM auth_user\n",
    "# \"\"\"\n",
    "# user_df = pd.read_sql_query(query, engine)\n",
    "# user_df.head()\n",
    "\n",
    "test_user_id = 4688272\n",
    "print(\n",
    "    clip_df[clip_df[\"user_id\"] == test_user_id][\"created_at\"].apply(lambda x: str(x)[:10]).value_counts()\n",
    ")\n",
    "print(clip_df[clip_df[\"user_id\"] == test_user_id].shape)\n",
    "query = \"\"\"\n",
    "SELECT *\n",
    "FROM auth_user\n",
    "WHERE id=27205089\n",
    "\"\"\"\n",
    "# 3 keenan\n",
    "# 6 martin\n",
    "# 8 tony -- that's me!\n",
    "# 186417 georg\n",
    "test_user_df = pd.read_sql_query(query, engine)\n",
    "test_user_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Find some weird generations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_bot_detection(clip_df, reaction_df, write_to_file=True, cut_off_freq=0.95)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Alpha testing user selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Alpha testing user selection\n",
    "# we focus on the folks who are good good\n",
    "\n",
    "# # 0526 is v2 -- prod\n",
    "# # 0529 is v4 -- still good IMO, more data\n",
    "# early_v3p5_data = pd.read_csv(\"/home/tony/Data/Preference/13b_v0/interesting_clips_20240529.csv\")\n",
    "\n",
    "# print(\"uqniue users for vp5\", early_v3p5_data[\"user_id\"].nunique())\n",
    "\n",
    "# early_v3_data = pd.read_csv(\"/home/tony/Data/Preference/7b_v0_interesting_clips.csv\")\n",
    "\n",
    "# print(\"uqniue users for v3\", early_v3_data[\"user_id\"].nunique())\n",
    "\n",
    "# early_v2_data = pd.read_csv(\"/home/tony/Data/Preference/3b_v0_interesting_clips.csv\")\n",
    "\n",
    "# print(\"uqniue users for v2\", early_v2_data[\"user_id\"].nunique())\n",
    "\n",
    "# intersection_user_ids_super = set(early_v3p5_data[\"user_id\"].unique()).intersection(set(early_v3_data[\"user_id\"].unique())).intersection(set(early_v2_data[\"user_id\"].unique()))\n",
    "\n",
    "# intersection_user_ids_v3_on = set(early_v3p5_data[\"user_id\"].unique()).intersection(set(early_v3_data[\"user_id\"].unique())).difference(intersection_user_ids_super)\n",
    "\n",
    "# print(len(intersection_user_ids_super), len(intersection_user_ids_v3_on))\n",
    "\n",
    "# super_user_df = user_df[user_df[\"id\"].isin(intersection_user_ids_super)].copy()\n",
    "# print(super_user_df.shape)\n",
    "# v3_onward_user_df = user_df[user_df[\"id\"].isin(intersection_user_ids_v3_on)].copy()\n",
    "# print(v3_onward_user_df.shape)\n",
    "# super_user_df.to_csv(\"/home/tony/Data/Preference/alpha_users/super_user.csv\", index=False)\n",
    "# v3_onward_user_df.to_csv(\"/home/tony/Data/Preference/alpha_users/v3_onward_user.csv\", index=False)\n",
    "# print(\"Done!!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# User generated clips lifetime filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"\"\"\n",
    "SELECT *\n",
    "FROM bots_userstats\n",
    "WHERE total_clips>=20\n",
    "\"\"\"\n",
    "user_stats_df = pd.read_sql_query(query, engine)\n",
    "print(user_stats_df.shape)\n",
    "user_stats_df[\"total_clips\"].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_users = user_stats_df[user_stats_df[\"total_clips\"] >= 100][\"user_id\"].unique()\n",
    "print(len(top_users))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-21T19:35:17.755108Z",
     "iopub.status.busy": "2024-06-21T19:35:17.754937Z",
     "iopub.status.idle": "2024-06-21T19:35:17.774581Z",
     "shell.execute_reply": "2024-06-21T19:35:17.774106Z",
     "shell.execute_reply.started": "2024-06-21T19:35:17.755091Z"
    }
   },
   "source": [
    "# Snow flake access"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(snow_password_path):\n",
    "    raise Exception(\"you are not authorized to access snowflake -- please setup\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "snow_session = Session.builder.configs(CONNECTION_PARAMETERS).create()\n",
    "\n",
    "snow_root = Root(snow_session)\n",
    "snow_schema = snow_root.databases[\"SUNO_PROD\"].schemas[\"PROD\"]\n",
    "print(snow_schema.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select the df we want to squery for play counts\n",
    "subset_v4_clips_df_all = final_interesting_clips[\n",
    "    final_interesting_clips[\"model_name\"] == target_model_name\n",
    "].copy()\n",
    "print(subset_v4_clips_df_all.shape)\n",
    "pair_request_mask = subset_v4_clips_df_all[\"request_id\"].isin(\n",
    "    subset_v4_clips_df_all[\"request_id\"]\n",
    "    .value_counts()\n",
    "    .index[subset_v4_clips_df_all[\"request_id\"].value_counts() == 2]\n",
    ")\n",
    "subset_v4_clips_df = subset_v4_clips_df_all[pair_request_mask].copy()\n",
    "print(subset_v4_clips_df.shape)\n",
    "v4_clip_ids = list(str(s) for s in subset_v4_clips_df[\"id\"].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "snow_batch_size = 100_000\n",
    "snow_results = []\n",
    "for clip_ids_chunk in tqdm.tqdm(\n",
    "    [v4_clip_ids[i : i + snow_batch_size] for i in range(0, len(v4_clip_ids), snow_batch_size)]\n",
    "):\n",
    "    id_query_str = \",\".join(\"'\" + x + \"'\" for x in clip_ids_chunk)\n",
    "    print(f\"Number of clip IDs in this chunk: {len(clip_ids_chunk)}\")\n",
    "    print(f\"Length of the ID query string: {len(id_query_str)}\")\n",
    "\n",
    "    session_query = snow_session.sql(\n",
    "        f\"\"\" select *\n",
    "        from ML_SONG_SUMMARY_INFO\n",
    "        where p_date = DATE(SYSDATE() - INTERVAL '2 HOUR')\n",
    "        and p_hour = hour(SYSDATE() - INTERVAL '2 HOUR')\n",
    "        and song_id in ({id_query_str})\n",
    "        order by p_hour desc;\"\"\"\n",
    "    )\n",
    "    temp_df_snow_test = pd.DataFrame(session_query.collect())\n",
    "    snow_results.append(temp_df_snow_test)\n",
    "print(len(snow_results))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_snow_test = pd.concat(snow_results)\n",
    "df_snow_test = df_snow_test.rename(columns=lambda x: x.lower())\n",
    "df_snow_test = df_snow_test.rename(columns={\"song_id\": \"str_id\"})\n",
    "print(\"Shape of df_snow_test:\")\n",
    "print(f\"Rows: {df_snow_test.shape[0]}\")\n",
    "print(f\"Columns: {df_snow_test.shape[1]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_v4_clips_df[\"str_id\"] = subset_v4_clips_df[\"id\"].astype(str)\n",
    "subset_v4_clips_df_test = subset_v4_clips_df.merge(df_snow_test, on=\"str_id\", how=\"left\")\n",
    "subset_v4_clips_df_test[\"norm_play_frac\"] = (\n",
    "    subset_v4_clips_df_test[\"total_play_time\"].fillna(0) / subset_v4_clips_df_test[\"duration\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a figure with two subplots\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))\n",
    "\n",
    "# First subplot: Total play duration\n",
    "pos_play_time = subset_v4_clips_df_test[subset_v4_clips_df_test[\"preference\"]][\"total_play_time\"]\n",
    "neg_play_time = subset_v4_clips_df_test[~subset_v4_clips_df_test[\"preference\"]][\"total_play_time\"]\n",
    "\n",
    "pos_play_time.hist(\n",
    "    bins=np.linspace(0, 400, 100),\n",
    "    alpha=0.5,\n",
    "    label=f\"pos (mean={pos_play_time.mean():.2f}, median={pos_play_time.median():.2f})\",\n",
    "    ax=ax1,\n",
    ")\n",
    "neg_play_time.hist(\n",
    "    bins=np.linspace(0, 400, 100),\n",
    "    alpha=0.5,\n",
    "    label=f\"neg (mean={neg_play_time.mean():.2f}, median={neg_play_time.median():.2f})\",\n",
    "    ax=ax1,\n",
    ")\n",
    "ax1.legend()\n",
    "ax1.set_xlabel(\"Total play duration in seconds\")\n",
    "ax1.set_ylabel(\"counts\")\n",
    "ax1.set_title(\"Play duration comparison\")\n",
    "\n",
    "# Second subplot: Normalized play fraction\n",
    "pos_norm_play_frac = subset_v4_clips_df_test[subset_v4_clips_df_test[\"preference\"]][\"norm_play_frac\"]\n",
    "neg_norm_play_frac = subset_v4_clips_df_test[~subset_v4_clips_df_test[\"preference\"]][\"norm_play_frac\"]\n",
    "\n",
    "pos_norm_play_frac.hist(\n",
    "    bins=np.linspace(0, 10, 100),\n",
    "    alpha=0.5,\n",
    "    label=f\"pos (mean={pos_norm_play_frac.mean():.2f}, median={pos_norm_play_frac.median():.2f})\",\n",
    "    ax=ax2,\n",
    ")\n",
    "neg_norm_play_frac.hist(\n",
    "    bins=np.linspace(0, 10, 100),\n",
    "    alpha=0.5,\n",
    "    label=f\"neg (mean={neg_norm_play_frac.mean():.2f}, median={neg_norm_play_frac.median():.2f})\",\n",
    "    ax=ax2,\n",
    ")\n",
    "ax2.legend()\n",
    "ax2.set_xlabel(\"Normalized play counts (play duration/duration)\")\n",
    "ax2.set_ylabel(\"Log counts\")\n",
    "ax2.set_yscale(\"log\")\n",
    "ax2.set_title(\"Normalized play duration comparison (Log scale)\")\n",
    "\n",
    "# Adjust layout and display the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "play_duration_mask = (\n",
    "    subset_v4_clips_df_test[\"preference\"]\n",
    "    & (subset_v4_clips_df_test[\"norm_play_frac\"] >= 0.95)\n",
    "    & (subset_v4_clips_df_test[\"total_play_time\"] >= 10)\n",
    "    & (subset_v4_clips_df_test[\"user_id\"].isin(top_users))\n",
    ") | (\n",
    "    (~subset_v4_clips_df_test[\"preference\"])\n",
    "    & (subset_v4_clips_df_test[\"norm_play_frac\"] <= 3.1)\n",
    "    & (subset_v4_clips_df_test[\"total_play_time\"] >= 10)\n",
    "    & (subset_v4_clips_df_test[\"user_id\"].isin(top_users))\n",
    ")\n",
    "# Calculate the fraction of clips that pass the play duration cut\n",
    "frac_pass_play_duration = play_duration_mask.sum() / subset_v4_clips_df_test.shape[0]\n",
    "\n",
    "# Print the result with a formatted string\n",
    "print(f\"Fraction of clips that pass the play duration cut: {frac_pass_play_duration:.4f}\")\n",
    "\n",
    "# Get unique request IDs that pass the play duration criteria\n",
    "unique_requests_pass_play_durations = subset_v4_clips_df_test[play_duration_mask][\"request_id\"].unique()\n",
    "\n",
    "# Print the number of unique requests that pass the play duration criteria\n",
    "print(\n",
    "    f\"Number of unique requests passing play duration criteria: {len(unique_requests_pass_play_durations)}\"\n",
    ")\n",
    "\n",
    "# Calculate the fraction of unique requests that pass play duration criteria\n",
    "fraction_requests_pass = (\n",
    "    len(unique_requests_pass_play_durations) / subset_v4_clips_df_test[\"request_id\"].nunique()\n",
    ")\n",
    "\n",
    "# Print the result with a formatted string\n",
    "print(f\"Fraction of unique requests that pass play duration criteria: {fraction_requests_pass:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_v4_clips_df_pass_duration = subset_v4_clips_df_test[play_duration_mask].copy()\n",
    "play_duration_mask_request_mask = subset_v4_clips_df_pass_duration[\"request_id\"].isin(\n",
    "    subset_v4_clips_df_pass_duration[\"request_id\"]\n",
    "    .value_counts()\n",
    "    .index[subset_v4_clips_df_pass_duration[\"request_id\"].value_counts() == 2]\n",
    ")\n",
    "final_subset_v4_clips_df = subset_v4_clips_df_pass_duration[play_duration_mask_request_mask].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count and print the number of unique request IDs\n",
    "unique_request_count = final_subset_v4_clips_df[\"request_id\"].nunique()\n",
    "print(\n",
    "    f\"Number of unique request IDs: {unique_request_count:,}, Total {subset_v4_clips_df_test['request_id'].nunique()}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_v4_clips_df_test[\"task\"].value_counts(), final_subset_v4_clips_df[\"task\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# final_subset_v4_clips_df.to_pickle(\n",
    "#     \"/home/tony/Data/Preference/30b_v2/interesting_clips_v4_t_3_20240912_full.pkl\",\n",
    "# )\n",
    "print(final_subset_v4_clips_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset_v4_clips_df.to_pickle(\n",
    "#     \"/home/tony/Data/Preference/30b_v2/interesting_clips_v4_t_3_20240901_full.pkl\",\n",
    "# )\n",
    "# print(subset_v4_clips_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset_v4_clips_df_test[subset_v4_clips_df_test[\"task\"] == \"infill\"][\"id_x\"] #[\"total_play_time\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_df[\"task\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
