data_processing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "\n",
    "# check what data sources we have\n",
    "# Working data sources:\n",
    "#  - Facebook data dump in JSON format (the whole folder named 'facebook-YourName12345' in data/)\n",
    "#  - Discord GDPR request data (the whole folder renamed to 'discord-YourUsername#1234' in data/)\n",
    "#  - Google Takeout dump of Hangouts (the whole folder renamed to 'google-YourUsername' in data/)\n",
    "dirs = os.listdir('data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# load in saved aliases\n",
    "with open('aliases.json') as f:\n",
    "    alias = json.load(f)\n",
    "    \n",
    "# set who \"me\" is for all this data\n",
    "me = 'cephcyn'\n",
    "\n",
    "# set the keyword that we use to denote \"media\"\n",
    "media_word = '[media_goes_here]'\n",
    "\n",
    "# set flag for if we want to save figures\n",
    "save_figures = True\n",
    "\n",
    "# translate from datasource name to unaliased name (if available)\n",
    "def un_alias(name):\n",
    "    if name in alias.keys():\n",
    "        return alias[name]\n",
    "    else:\n",
    "        return name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize data formats\n",
    "from collections import namedtuple\n",
    "\n",
    "Message = namedtuple('Message', ['person', 'sent_by_me', 'text', 'timestamp', 'valid'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clear the log\n",
    "\n",
    "log = []\n",
    "with open('dumps/msgs.pickle', 'wb') as handle:\n",
    "    pickle.dump(log, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read Facebook Messenger data into custom format\n",
    "import re\n",
    "import json\n",
    "from datetime import datetime, timezone\n",
    "\n",
    "fb_dir_name_rx = re.compile(\"^facebook-.*$\")\n",
    "json_rx = re.compile(\"^.*\\.json$\")\n",
    "\n",
    "def parse_msg(msg_data):\n",
    "    if len(msg_data['participants']) != 2:\n",
    "        return []\n",
    "    \n",
    "    # Sanity check that I'm in the chat\n",
    "    # also get the name of the person I'm talking with\n",
    "    participants = [un_alias(p['name']) for p in msg_data['participants']]\n",
    "    if (participants[0] == me):\n",
    "        conversant = participants[1]\n",
    "    elif (participants[1] == me):\n",
    "        conversant = participants[0]\n",
    "    else:\n",
    "        print('\\'me\\' not found in chat, add alias?')\n",
    "        print('Participants:', msg_data['participants'])\n",
    "        return []\n",
    "    \n",
    "    # Exclude the conversant if null (alias file says to remove them)\n",
    "    if conversant is None:\n",
    "        return []\n",
    "    \n",
    "    # Parse the chat log into a list of tuples\n",
    "    log = []\n",
    "    for msg in msg_data['messages']:\n",
    "        # note that Facebook saves timestamp_mS not S\n",
    "        msg['sender_name'] = un_alias(msg['sender_name'])\n",
    "        ms = msg['timestamp_ms']\n",
    "        try:\n",
    "            msg_tuple = Message(\n",
    "                conversant, \n",
    "                msg['sender_name'] == me, \n",
    "                msg['content'],\n",
    "                datetime.fromtimestamp(ms//1000, tz=timezone.utc).replace(microsecond=ms%1000*1000),\n",
    "                True\n",
    "            )\n",
    "        except:\n",
    "            # if there's no content, assume it's a sticker or photo or video or whatever\n",
    "            msg_tuple = Message(\n",
    "                conversant, \n",
    "                msg['sender_name'] == me, \n",
    "                media_word,\n",
    "                datetime.fromtimestamp(ms//1000, tz=timezone.utc).replace(microsecond=ms%1000*1000),\n",
    "                True\n",
    "            )\n",
    "        log.append(msg_tuple)\n",
    "    return log\n",
    "\n",
    "def parse_dir(fb_dir_name, msg_type_dir):\n",
    "    log = []\n",
    "    for chat_name in os.listdir(f'data/{fb_dir_name}/messages/{msg_type_dir}'):\n",
    "        for chat_fname in os.listdir(f'data/{fb_dir_name}/messages/{msg_type_dir}/{chat_name}'):\n",
    "            if json_rx.match(chat_fname):\n",
    "                with open(f'data/{fb_dir_name}/messages/{msg_type_dir}/{chat_name}/{chat_fname}') as f:\n",
    "                    msg_data = json.load(f)\n",
    "                    log = log + parse_msg(msg_data)\n",
    "    return log\n",
    "\n",
    "# Load already-existing logs!\n",
    "with open('dumps/msgs.pickle', 'rb') as handle:\n",
    "    log = pickle.load(handle)\n",
    "\n",
    "for fb_dir_name in dirs:\n",
    "    if fb_dir_name_rx.match(fb_dir_name):\n",
    "        msg_dir = os.listdir(f'data/{fb_dir_name}/messages')\n",
    "        for msg_type_dir in ['inbox', 'archived_threads']:\n",
    "            if msg_type_dir in msg_dir:\n",
    "                log = log + parse_dir(fb_dir_name, msg_type_dir)\n",
    "\n",
    "# Save log with Facebook data to disk\n",
    "with open('dumps/msgs.pickle', 'wb') as handle:\n",
    "    pickle.dump(log, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read Discord data dump (directly obtained from Discord) into custom format\n",
    "# note that this parse will not give us names for people who aren't added as friends on Discord\n",
    "import re\n",
    "import json\n",
    "import csv\n",
    "from datetime import datetime\n",
    "\n",
    "dc_dir_name_rx = re.compile('^discord-.*$')\n",
    "msg_log_dir_rx = re.compile('^[0-9]+$')\n",
    "\n",
    "def parse_msg(id_to_dc_uname, msg_log_dir, show_error_messages=True):\n",
    "    with open(f'data/{dc_dir_name}/messages/{msg_log_dir}/channel.json') as f:\n",
    "        channel_info = json.load(f)\n",
    "    \n",
    "    # skip over chats that are group chats or server channels\n",
    "    if not channel_info['type'] == 1:\n",
    "        return []\n",
    "\n",
    "    # Sanity check that I'm in the chat\n",
    "    # also get the name of the person that I'm chatting with\n",
    "    try:\n",
    "        participants = [un_alias(id_to_dc_uname[p]) for p in channel_info['recipients']]\n",
    "    except:\n",
    "        if show_error_messages:\n",
    "            print('at least one of the Discord DM participants is unknown (not a friend?)')\n",
    "            print('Participant IDs:    ', channel_info['recipients'])\n",
    "        return []\n",
    "    if (participants[0] == me):\n",
    "        conversant = participants[1]\n",
    "    elif (participants[1] == me):\n",
    "        conversant = participants[0]\n",
    "    else:\n",
    "        if show_error_messages:\n",
    "            print('\\'me\\' not found in chat, add alias?')\n",
    "            print('Participant handles:', participants)\n",
    "        return []\n",
    "    \n",
    "    # Exclude the conversant if null (alias file says to remove them)\n",
    "    if conversant is None:\n",
    "        return []\n",
    "    \n",
    "    # Parse the chat log into a list of tuples\n",
    "    log = []\n",
    "    with open(f'data/{dc_dir_name}/messages/{msg_log_dir}/messages.csv') as f:\n",
    "        # note that Discord data dump only includes messages that I sent, not both-party messages\n",
    "        csv_reader = csv.DictReader(f)\n",
    "        for row in csv_reader:\n",
    "            try:\n",
    "                msg_time = datetime.strptime(row['Timestamp'], '%Y-%m-%d %H:%M:%S.%f%z')\n",
    "            except:\n",
    "                # Discord exports data format without ms sometimes??? not sure why\n",
    "                msg_time = datetime.strptime(row['Timestamp'], '%Y-%m-%d %H:%M:%S%z')\n",
    "            msg_tuple = Message(\n",
    "                conversant,\n",
    "                True, # Discord logs only save messages sent by me\n",
    "                row['Contents'],\n",
    "                msg_time,\n",
    "                True\n",
    "            )\n",
    "            log.append(msg_tuple)\n",
    "            # For basic graph sanity's sake, assume that for each message I sent on Discord, I got one back\n",
    "            # this definitely isn't accurate, but it's the best I can do with the official data dump\n",
    "            msg_tuple = Message(\n",
    "                conversant,\n",
    "                False,\n",
    "                row['Contents'],\n",
    "                msg_time,\n",
    "                False # this is a BS message, we're pretending ppl always reply and same-day\n",
    "            )\n",
    "            log.append(msg_tuple)\n",
    "    return log\n",
    "\n",
    "# Load already-existing logs!\n",
    "with open('dumps/msgs.pickle', 'rb') as handle:\n",
    "    log = pickle.load(handle)\n",
    "\n",
    "for dc_dir_name in dirs:\n",
    "    if dc_dir_name_rx.match(dc_dir_name):\n",
    "        # Get the friend username data\n",
    "        # maps {id: username#discriminator} i.e. {3243728943728:\"Username#0001\"}\n",
    "        id_to_dc_uname = {}\n",
    "        with open(f'data/{dc_dir_name}/account/user.json') as f:\n",
    "            dc_user_data = json.load(f)\n",
    "            # load in my ID\n",
    "            my_username = dc_user_data['username']\n",
    "            my_discriminator = dc_user_data['discriminator']\n",
    "            id_to_dc_uname[dc_user_data['id']] = f'{my_username}#{my_discriminator}'\n",
    "            # load in friends' IDs\n",
    "            for friend in dc_user_data['relationships']:\n",
    "                friend_username = friend['user']['username']\n",
    "                friend_discriminator = friend['user']['discriminator']\n",
    "                id_to_dc_uname[friend['id']] = f'{friend_username}#{friend_discriminator}'\n",
    "        \n",
    "        # Get the sent-messages data\n",
    "        msg_dir = os.listdir(f'data/{dc_dir_name}/messages')\n",
    "        for msg_log_dir in msg_dir:\n",
    "            if msg_log_dir_rx.match(msg_log_dir):\n",
    "                log = log + parse_msg(id_to_dc_uname, msg_log_dir)\n",
    "        \n",
    "# Save log with Discord dump data to disk\n",
    "with open('dumps/msgs.pickle', 'wb') as handle:\n",
    "    pickle.dump(log, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read Discord data scrape (using DiscordChatExporter) into custom format\n",
    "# TODO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read Google Takeout data export of Hangouts chat logs into custom format\n",
    "import re\n",
    "import json\n",
    "import csv\n",
    "from datetime import datetime, timezone\n",
    "\n",
    "gg_dir_name_rx = re.compile('^google-.*$')\n",
    "\n",
    "# Load already-existing logs!\n",
    "with open('dumps/msgs.pickle', 'rb') as handle:\n",
    "    log = pickle.load(handle)\n",
    "\n",
    "def parse_msg(msg_data, id_to_gg_uname):\n",
    "    # TODO add support for 2-person group chat data?\n",
    "    if not msg_data['conversation']['conversation']['type'] == 'STICKY_ONE_TO_ONE':\n",
    "        return []\n",
    "    if len(msg_data['conversation']['conversation']['participant_data']) > 2:\n",
    "        return []\n",
    "    \n",
    "    # Sanity check that I'm in the chat\n",
    "    participants = msg_data['conversation']['conversation']['participant_data']\n",
    "    participants = [un_alias(id_to_gg_uname[p['id']['gaia_id']]) for p in participants]\n",
    "    if (participants[0] == me):\n",
    "        conversant = participants[1]\n",
    "    elif (participants[1] == me):\n",
    "        conversant = participants[0]\n",
    "    else:\n",
    "        print('\\'me\\' not found in chat, add alias?')\n",
    "        print('Participants:', msg_data['conversation']['conversation']['participant_data'])\n",
    "        return []\n",
    "    \n",
    "    # Exclude the conversant if null (alias file says to remove them)\n",
    "    if conversant is None:\n",
    "        return []\n",
    "    \n",
    "    # Parse the chat log into a list of tuples\n",
    "    log = []\n",
    "    for msg in msg_data['events']:\n",
    "        msg['sender_name'] = un_alias(id_to_gg_uname[msg['sender_id']['gaia_id']])\n",
    "        # note that Google saves timestamp_uS, not S\n",
    "        ms = int(msg['timestamp'])\n",
    "        try:\n",
    "            msg_tuple = Message(\n",
    "                conversant, \n",
    "                msg['sender_name'] == me, \n",
    "                ' '.join([s['text'] for s in msg['chat_message']['message_content']['segment']]),\n",
    "                datetime.fromtimestamp(ms//1000000, tz=timezone.utc).replace(microsecond=ms%1000000),\n",
    "                True\n",
    "            )\n",
    "        except:\n",
    "            # if there's no content, assume it's a photo or video or Hangouts call or whatever\n",
    "            msg_tuple = Message(\n",
    "                conversant, \n",
    "                msg['sender_name'] == me, \n",
    "                media_word,\n",
    "                datetime.fromtimestamp(ms//1000000, tz=timezone.utc).replace(microsecond=ms%1000000),\n",
    "                True\n",
    "            )\n",
    "        log.append(msg_tuple)\n",
    "    return log\n",
    "\n",
    "for gg_dir_name in dirs:\n",
    "    if gg_dir_name_rx.match(gg_dir_name):\n",
    "        with open(f'data/{gg_dir_name}/Hangouts/Hangouts.json') as f:\n",
    "            gg_user_data = json.load(f)\n",
    "        \n",
    "        # Build up the mapping from gaia_id->name\n",
    "        id_to_gg_uname = {}\n",
    "        for i in range(len(gg_user_data['conversations'])):\n",
    "            msg_data = gg_user_data['conversations'][i]\n",
    "            participants = msg_data['conversation']['conversation']['participant_data']\n",
    "            # TODO double check if we should use chat_id or gaia_id?\n",
    "            for p in participants:\n",
    "                # only add a mapping if we do not have a proper mapping yet\n",
    "                if (p['id']['gaia_id'] not in id_to_gg_uname.keys()) \\\n",
    "                        or (id_to_gg_uname[p['id']['gaia_id']] is None):\n",
    "                    if 'fallback_name' in p.keys():\n",
    "                        id_to_gg_uname[p['id']['gaia_id']] = p['fallback_name']\n",
    "                    else:\n",
    "                        id_to_gg_uname[p['id']['gaia_id']] = None\n",
    "        \n",
    "        # Parse messages\n",
    "        for i in range(len(gg_user_data['conversations'])):\n",
    "            log = log + parse_msg(gg_user_data['conversations'][i], id_to_gg_uname)\n",
    "        \n",
    "# # Save log with Google dump data to disk\n",
    "with open('dumps/msgs.pickle', 'wb') as handle:\n",
    "    pickle.dump(log, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "# Load the data!\n",
    "with open('dumps/msgs.pickle', 'rb') as handle:\n",
    "    messages = pickle.load(handle)\n",
    "print(\"Total Messages:\", len(messages))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See the list of people included\n",
    "print(sorted(set([m[0] for m in messages])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, timedelta\n",
    "import pytz\n",
    "\n",
    "# Set bounds for data structures for the graphs\n",
    "start_date = datetime.strptime('3/27/13 +00:00', \"%m/%d/%y %z\")\n",
    "end_date = datetime.strptime('5/27/20 +00:00', \"%m/%d/%y %z\")\n",
    "n_days = (end_date - start_date).days\n",
    "n_lengths = 10000\n",
    "\n",
    "# Set the granularity for time-based graphs\n",
    "delta_timechunks = timedelta(minutes=5)\n",
    "n_timechunks = timedelta(days=1) // delta_timechunks\n",
    "\n",
    "# Set your local timezone\n",
    "local_tz = pytz.timezone('Etc/GMT-7')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "import numpy as np\n",
    "from collections import OrderedDict\n",
    "import pandas as pd\n",
    "\n",
    "person_start_date = {} # person_start_date[name] = datetime of first contact\n",
    "msgs_by_day = {} # msgs_by_day[name] = [] * n_days, populated with message counts\n",
    "msgs_by_time = {} # msgs_by_time[name] = [] * n_timechunks, populated with message counts\n",
    "msgs_by_length = {} # msgs_by_length[name] = [] * n_lengths, populated with message counts\n",
    "\n",
    "for msg in messages:\n",
    "    # fill in person_start_date\n",
    "    if msg.person in person_start_date.keys():\n",
    "        if msg.timestamp < person_start_date[msg.person]:\n",
    "            person_start_date[msg.person] = msg.timestamp\n",
    "    else:\n",
    "        person_start_date[msg.person] = msg.timestamp\n",
    "    # fill in msgs_by_day\n",
    "    if msg.person in msgs_by_day.keys():\n",
    "        if start_date < msg.timestamp < end_date:\n",
    "                idx = (msg.timestamp - start_date).days\n",
    "                msgs_by_day[msg.person][idx] += 1\n",
    "    else:\n",
    "        msgs_by_day[msg.person] = [0] * n_days\n",
    "    # fill in msgs_by_time\n",
    "    if msg.person in msgs_by_time.keys():\n",
    "        if start_date < msg.timestamp < end_date:\n",
    "                localtime = msg.timestamp.astimezone()\n",
    "                idx = (localtime - localtime.replace(hour=0, minute=0, second=0, microsecond=0)) // delta_timechunks\n",
    "                msgs_by_time[msg.person][idx] += 1\n",
    "    else:\n",
    "        msgs_by_time[msg.person] = [0] * n_timechunks\n",
    "    # fill in msgs_by_length\n",
    "    if msg.person in msgs_by_length.keys():\n",
    "        if start_date < msg.timestamp < end_date:\n",
    "            msgs_by_length[msg.person][len(msg.text)] += 1\n",
    "    else:\n",
    "        msgs_by_length[msg.person] = [0] * n_lengths\n",
    "        \n",
    "msgs_by_day = OrderedDict(sorted(msgs_by_day.items(), key=lambda i: -sum(i[1])))\n",
    "msgs_by_time = OrderedDict(sorted(msgs_by_time.items(), key=lambda i: -sum(i[1])))\n",
    "msgs_by_length = OrderedDict(sorted(msgs_by_length.items(), key=lambda i: -sum(i[1])))\n",
    "\n",
    "msg_count_data = pd.DataFrame(msgs_by_day, index=pd.date_range(start_date, periods=n_days))\n",
    "msg_time_data = pd.DataFrame(msgs_by_time, index=pd.timedelta_range(start='0', periods=n_timechunks, freq=delta_timechunks))\n",
    "msg_length_data = pd.DataFrame(msgs_by_length)\n",
    "\n",
    "msg_datasets = [msg_count_data, msg_time_data, msg_length_data]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, value in sorted(person_start_date.items(), key=lambda i: i[1]):\n",
    "    print(key, '\\t: ', value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dropout any people we don't want to include in the results\n",
    "# for d in msg_datasets:\n",
    "#     del d['SomeName1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def numTalkedToPlot(data, min_messages=1, rolling_window=1):\n",
    "    talkedTo = data[ min_messages < data ]\n",
    "    talkedTo = ~pd.isnull(talkedTo)\n",
    "    toPlot = talkedTo.iloc[:,:].sum(axis=1)\n",
    "    return toPlot.rolling(rolling_window).mean()\n",
    "\n",
    "plt.clf()\n",
    "toPlot = numTalkedToPlot(msg_count_data, rolling_window=7)\n",
    "ax = toPlot.plot(title=\"Number of People Talked to\", figsize=(10,3))\n",
    "ax.set_ylabel(\"Number of People\")\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(\"outputs/people_count_dates.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.clf()\n",
    "toPlot = numTalkedToPlot(msg_time_data, rolling_window=3)\n",
    "ax = toPlot.plot(title=\"Number of People ever Talked To At Time\", figsize=(10,3))\n",
    "ax.set_ylabel(\"Number of People\")\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(\"outputs/people_count_times.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def everyonePlot(data, rolling_window=1):\n",
    "    sum_data = data.iloc[:,:].sum(axis=1)\n",
    "    return sum_data.rolling(rolling_window).mean()\n",
    "\n",
    "plt.clf()\n",
    "toPlot = everyonePlot(msg_count_data, rolling_window=1)\n",
    "ax = toPlot.plot(title=\"Message Count Data\", figsize=(10,3))\n",
    "ax.set_ylabel(\"Number of Messages\")\n",
    "ax.set_xlabel(\"Date\")\n",
    "ax.set_ylim((0, toPlot.max()*1.1))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(\"outputs/messages_count_dates.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.clf()\n",
    "toPlot = everyonePlot(msg_time_data, rolling_window=3)\n",
    "ax = toPlot.plot(title=\"Message Timestamp Data\", figsize=(10,3))\n",
    "ax.set_ylabel(\"Number of Messages\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "ax.set_ylim((0, toPlot.max()*1.1))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(\"outputs/messages_count_times.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.clf()\n",
    "toPlot = everyonePlot(msg_length_data, rolling_window=1)\n",
    "ax = toPlot.plot(title=\"Message Length Data\", figsize=(10,3))\n",
    "ax.set_ylabel(\"Number of Messages\")\n",
    "ax.set_xlabel(\"Length\")\n",
    "ax.set_ylim((0, toPlot.max()*1.1))\n",
    "ax.set_xlim((0, 100))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(\"outputs/messages_length.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cumMsgPlot(data, start, end):\n",
    "    data = data.copy().cumsum(axis=0)\n",
    "    data = data.iloc[:, start:end]\n",
    "    return data\n",
    "\n",
    "plt.clf()\n",
    "top_n = 6\n",
    "toPlot = cumMsgPlot(msg_count_data, start=0, end=top_n)\n",
    "ax = toPlot.plot(title=f\"Cumulative Messaging Data (Top {top_n} Most Talked To)\", legend=True, figsize=(10,7))\n",
    "ax.set_ylabel(\"Cumulative Number of Messages\")\n",
    "ax.set_xlabel(\"Date\")\n",
    "plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(f\"outputs/messages_count_cumulative_top{top_n}.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def avgMsgPlot(data, start, end, rolling_window=1):\n",
    "#     data = data.copy().cumsum(axis=0)\n",
    "    data = data.iloc[:, start:end]\n",
    "    return data.rolling(rolling_window).mean()\n",
    "\n",
    "plt.clf()\n",
    "top_n = 6\n",
    "toPlot = avgMsgPlot(msg_count_data, start=0, end=top_n, rolling_window=30)\n",
    "ax = toPlot.plot(title=f\"Active Messaging Data (Top {top_n} Most Talked To)\", legend=True, figsize=(10,7))\n",
    "ax.set_ylabel(\"Avg Number of Messages per Day\")\n",
    "ax.set_xlabel(\"Date\")\n",
    "plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(f\"outputs/messages_count_active_top{top_n}.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def indivMsgPlotTopN(data, start, end, rolling_window=12):\n",
    "    data = data.iloc[:, start:end]\n",
    "    return data.rolling(rolling_window).mean()\n",
    "\n",
    "plt.clf()\n",
    "top_n = 6\n",
    "toPlot = indivMsgPlotTopN(msg_time_data, start=0, end=top_n)\n",
    "ax = toPlot.plot(title=f\"Message Time Data (Top {top_n} Most Talked To)\", legend=True, figsize=(10,7))\n",
    "ax.set_ylabel(\"Cumulative Number of Messages\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(f\"outputs/messages_count_times_top{top_n}.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def indivMsgPlot(data, select, rolling_window=12):\n",
    "    select = [data.columns.get_loc(x) for x in select]\n",
    "    data = data.iloc[:, select]\n",
    "    return data.rolling(rolling_window).mean()\n",
    "\n",
    "plt.clf()\n",
    "select = ['SomeName1', 'SomeName2', 'SomeName3', 'SomeName4']\n",
    "toPlot = indivMsgPlot(msg_time_data, select=select)\n",
    "ax = toPlot.plot(title=f\"Message Time Data\", legend=True, figsize=(10,7))\n",
    "ax.set_ylabel(\"Cumulative Number of Messages\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(f\"outputs/messages_count_times_selected.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def indivMsgLengthPlot(data, select, rolling_window=1):\n",
    "    select = [data.columns.get_loc(x) for x in select]\n",
    "    data = data.iloc[:, select]\n",
    "    data = data.iloc[:data.max().max(), :]\n",
    "    return data.rolling(rolling_window).mean()\n",
    "\n",
    "plt.clf()\n",
    "select = ['SomeName1', 'SomeName2', 'SomeName3', 'SomeName4']\n",
    "toPlot = indivMsgLengthPlot(msg_length_data, select=select, rolling_window=1)\n",
    "ax = toPlot.plot(title=f\"Message Length Data\", legend=True, figsize=(10,7))\n",
    "ax.set_ylabel(\"Cumulative Number of Messages\")\n",
    "ax.set_xlabel(\"Message Length\")\n",
    "ax.set_xlim((0, 100))\n",
    "plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(f\"outputs/messages_length_selected.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalizedIndivMsgPlot(data, select, rolling_window=12):\n",
    "    select = [data.columns.get_loc(x) for x in select]\n",
    "    data = data.iloc[:, select]\n",
    "    data = data.rolling(rolling_window).mean()\n",
    "    data = data.apply(lambda x: x/x.sum(), axis=0)\n",
    "    return data\n",
    "\n",
    "plt.clf()\n",
    "select = ['SomeName1', 'SomeName2', 'SomeName3', 'SomeName4']\n",
    "toPlot = normalizedIndivMsgPlot(msg_time_data, select=select)\n",
    "ax = toPlot.plot(title=f\"Message Time Data (Normalized)\", legend=True, figsize=(10,7))\n",
    "ax.set_ylabel(\"Message Count (Normalized)\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "fig = ax.get_figure()\n",
    "if save_figures:\n",
    "    fig.savefig(f\"outputs/messages_normcount_times_selected.png\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nonfunctional - we don't have complete data for messages sent/received\n",
    "#                 so this graph is not as important\n",
    "\n",
    "# messages_in_range = filter(lambda x: start_date < x.timestamp < end_date, messages)\n",
    "\n",
    "# df = pd.DataFrame(list(messages_in_range), columns=Message._fields)\n",
    "# df = df[[\"person\", \"sent_by_me\"]]\n",
    "# df.columns = [\"person\", \"sent\"]\n",
    "\n",
    "# df['received'] = pd.Series(~df[\"sent\"], index=df.index)\n",
    "# df['total'] = df['sent'] | df['received']\n",
    "\n",
    "# grouped = df.groupby('person')\n",
    "# sent_received = grouped.sum().sort_values('total', ascending=False)\n",
    "# toPlot = sent_received[[\"sent\", \"received\"]].ix[:15, :]\n",
    "\n",
    "# ax = toPlot.plot.bar(title=\"Total Messages Sent/Received (Top 15 Most Talked To)\", stacked=True, color=('b', 'r'), figsize=(10, 5))\n",
    "# ax.set_ylabel(\"Number of Messages\")\n",
    "# ax.set_xlabel(\"Person\")\n",
    "\n",
    "# fig = ax.get_figure()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat Topic Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess data for chat topic analysis\n",
    "import gensim\n",
    "from gensim.utils import simple_preprocess\n",
    "from gensim.parsing.preprocessing import STOPWORDS\n",
    "from nltk.stem import WordNetLemmatizer, SnowballStemmer\n",
    "from nltk.stem.porter import *\n",
    "import nltk\n",
    "nltk.download('wordnet')\n",
    "\n",
    "stemmer = SnowballStemmer(\"english\")\n",
    "\n",
    "msgs_per_person = {} # msgs_per_person[name] = [[all words they sent me], [all words i sent them]]\n",
    "\n",
    "def lemmatize_stemming(text):\n",
    "    return stemmer.stem(WordNetLemmatizer().lemmatize(text, pos='v'))\n",
    "\n",
    "# Tokenize and lemmatize\n",
    "def preprocess(text):\n",
    "    result = []\n",
    "    for token in gensim.utils.simple_preprocess(text) :\n",
    "        if token not in gensim.parsing.preprocessing.STOPWORDS:# and len(token) > 3:\n",
    "#             result.append(token)\n",
    "            result.append(lemmatize_stemming(token))\n",
    "    return result\n",
    "\n",
    "for msg in messages:\n",
    "    if msg.valid:\n",
    "        if msg.person in msgs_per_person.keys():\n",
    "            msgs_per_person[msg.person][msg.sent_by_me] += preprocess(msg.text)\n",
    "        else:\n",
    "            msgs_per_person[msg.person] = [[], []]\n",
    "            msgs_per_person[msg.person][msg.sent_by_me] += preprocess(msg.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dictionary = gensim.corpora.Dictionary(\n",
    "    [msgs_per_person[k][0] + msgs_per_person[k][0] for k in msgs_per_person.keys()]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gensim import models\n",
    "\n",
    "# Create bag-of-words representation of the chats\n",
    "corpus = [dictionary.doc2bow(\n",
    "    msgs_per_person[person][0] + msgs_per_person[person][1]\n",
    ") for person in msgs_per_person.keys()]\n",
    "# for doc in corpus:\n",
    "#     print([[dictionary[id], np.around(freq, decimals=2)] for id, freq in doc])\n",
    "\n",
    "# Create TF-IDF (term freq, inverse document freq) representation of the chats\n",
    "tfidf = models.TfidfModel(corpus)\n",
    "# for doc in tfidf[corpus]:\n",
    "#     print([[dictionary[id], np.around(freq, decimals=2)] for id, freq in doc])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import altair as alt\n",
    "alt.renderers.enable('default')\n",
    "alt.data_transformers.enable('json')\n",
    "\n",
    "summed_up_corpus = [0]*len(dictionary)\n",
    "word_counts = {}\n",
    "for doc in corpus:\n",
    "    for (word_id, count) in doc:\n",
    "        summed_up_corpus[word_id] += count\n",
    "for i in range(len(summed_up_corpus)):\n",
    "    word_counts[dictionary[i]] = summed_up_corpus[i]\n",
    "\n",
    "word_counts = pd.DataFrame(word_counts.items(), columns=['word', 'count'])\n",
    "word_counts = word_counts.sort_values(by=['count'], ascending=False)[:500]\n",
    "\n",
    "alt.Chart(word_counts).mark_bar().encode(\n",
    "    x=alt.X('word:N', sort='-y'),\n",
    "    y=alt.Y('count:Q', sort='ascending'),\n",
    "    tooltip=['word']\n",
    ")#.interactive()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getUniqueWordsTopN(data, name, top_n=50):\n",
    "    top_n_words = sorted(data[name], key=lambda x: x[1], reverse=True)[:top_n]\n",
    "    top_n_words = [(dictionary[id], freq) for id, freq in top_n_words]\n",
    "    return top_n_words\n",
    "\n",
    "print(getUniqueWordsTopN(dict(zip(msgs_per_person.keys(), tfidf[corpus])), 'SomeName'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create clusters for chats\n",
    "# Shamelessly adapted from another of my NLP projects\n",
    "# TODO: link to that, annotate this code so it is more understandable?\n",
    "\n",
    "from gensim.models import LdaModel\n",
    "\n",
    "def get_topics(dictionary, corpus, ids):\n",
    "    # Train LDA model, get model & topic vectors\n",
    "    # Set training parameters.\n",
    "    num_topics = 30\n",
    "    chunksize = 100\n",
    "    passes = 20\n",
    "    iterations = 400\n",
    "    eval_every = 100 # None # = Don't evaluate model perplexity, takes too much time.\n",
    "\n",
    "    # Make a index to word dictionary.\n",
    "    temp = dictionary[0]  # This is to \"load\" the dictionary.\n",
    "    id2word = dictionary.id2token\n",
    "\n",
    "    model = LdaModel(\n",
    "        corpus=corpus,\n",
    "        id2word=id2word,\n",
    "        chunksize=chunksize,\n",
    "        alpha='auto',\n",
    "        eta='auto',\n",
    "        iterations=iterations,\n",
    "        num_topics=num_topics,\n",
    "        passes=passes,\n",
    "        eval_every=eval_every\n",
    "    )\n",
    "    \n",
    "    # Get basic evaluation\n",
    "    top_topics = model.top_topics(corpus) #, num_words=20)\n",
    "\n",
    "    # Average topic coherence is the sum of topic coherences of all topics, divided by the number of topics.\n",
    "    avg_topic_coherence = sum([t[1] for t in top_topics]) / num_topics\n",
    "    print('Average topic coherence: %.4f.' % avg_topic_coherence)\n",
    "    \n",
    "    # Get topic vectors\n",
    "    all_topics = model.get_document_topics(corpus, per_word_topics=True)\n",
    "    all_topics = [(doc_topics, word_topics, word_phis) for doc_topics, word_topics, word_phis in all_topics]\n",
    "    sen_top = {}\n",
    "    for i in range(len(corpus)):\n",
    "        # These are in the same order as the documents themselves.\n",
    "        doc_topics, word_topics, phi_values = all_topics[i]\n",
    "        # Generate the topic VECTOR not just list of topics\n",
    "        doc_topic_vector = [0] * num_topics\n",
    "        for topic in doc_topics:\n",
    "            doc_topic_vector[topic[0]] = topic[1]\n",
    "        sen_top.update({ids[i]:doc_topic_vector})\n",
    "    \n",
    "    return model, sen_top\n",
    "\n",
    "model_bow, sen_top_bow = get_topics(dictionary, corpus, list(msgs_per_person.keys()))\n",
    "model_tfidf, sen_top_tfidf = get_topics(dictionary, tfidf[corpus], list(msgs_per_person.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the topic vectors\n",
    "\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.decomposition import PCA\n",
    "import altair as alt\n",
    "alt.renderers.enable('default')\n",
    "\n",
    "def visualize_embeds(data):\n",
    "    x = pd.DataFrame(data.values())\n",
    "    # x = StandardScaler().fit_transform(x)\n",
    "\n",
    "    pca = PCA(n_components=2)\n",
    "    principalComponents = pca.fit_transform(x)\n",
    "    principalDf = pd.DataFrame(data=principalComponents,\n",
    "                               columns = ['pc1', 'pc2'])\n",
    "\n",
    "    finalDf = pd.concat([principalDf, \n",
    "                         pd.DataFrame(sen_top_bow.keys()).rename(columns={0:'name'})],\n",
    "                        axis = 1)\n",
    "\n",
    "    return alt.Chart(finalDf).mark_circle(size=60).encode(\n",
    "        x='pc1',\n",
    "        y='pc2',\n",
    "#         color='type',\n",
    "        tooltip=['name']#, 'note', 'coord']\n",
    "    ).interactive()\n",
    "\n",
    "# visualize_embeds(sen_top_bow)\n",
    "visualize_embeds(sen_top_tfidf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cluster based on our topic vectors per-chat\n",
    "# Shamelessly adapted from another of my NLP projects\n",
    "# TODO: link to that, annotate this code so it is more understandable?\n",
    "\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "def similarity_clustering(similarity_dict, m, n):\n",
    "    clusters = {};\n",
    "    unselected_posts = similarity_dict.copy()\n",
    "    post_keys = list(unselected_posts.keys())\n",
    "    unselected_keys = list(unselected_posts.keys())\n",
    "    cluster_size = int(np.ceil(n / m))\n",
    "    # print(cluster_size)\n",
    "    while len(unselected_posts) != 0:\n",
    "        selected_post = random.choice(unselected_keys)\n",
    "        # labeling the selected row\n",
    "        emb_dict = dict(zip(post_keys, unselected_posts[selected_post]))\n",
    "        # only sort the unselected columns\n",
    "        sim = {k: emb_dict[k] for k in unselected_keys}\n",
    "        sim_sort = [k for k in sorted(sim.items(), key=lambda item: item[1])][::-1]\n",
    "        cluster_size = int(np.ceil(n / m))\n",
    "        try:\n",
    "            sim_most = sim_sort[0:cluster_size]\n",
    "        except:\n",
    "            sim_most = sim_sort[0:end]\n",
    "        clusters[selected_post] = sim_most\n",
    "        # deleted the selected rows from the unselected\n",
    "        for p in sim_most:\n",
    "            del unselected_posts[p[0]]\n",
    "        unselected_keys = list(unselected_posts.keys())\n",
    "        # print(cluster_size)\n",
    "    return clusters\n",
    "\n",
    "def cluster(numClusters, sen_emb):\n",
    "    numTotalPosts = len(msgs_per_person)\n",
    "    d = pd.DataFrame(sen_emb).transpose()\n",
    "    sim_mat = cosine_similarity(d)\n",
    "\n",
    "    post = list(d.index)\n",
    "    post_emb = dict(zip(post, sim_mat))\n",
    "\n",
    "    cluster = similarity_clustering(post_emb, numClusters, numTotalPosts)\n",
    "    print(cluster)\n",
    "\n",
    "    # Transform clusters into a name:cluster_id dict\n",
    "    transformed_cluster = {}\n",
    "    clust_num = 0\n",
    "    for key in cluster.keys():\n",
    "        for name in cluster[key]:\n",
    "            transformed_cluster[name[0]] = clust_num\n",
    "        clust_num += 1\n",
    "    print(transformed_cluster)\n",
    "    \n",
    "    return\n",
    "\n",
    "cluster(10, sen_top_bow)\n",
    "cluster(10, sen_top_tfidf)\n",
    "\n",
    "# TODO: incomplete - maybe visualize these somehow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}