Gmail_Mbox_to_TXT_JSON.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Prerequisite for This Program: A Mbox File with Emails\n",
        "\n",
        "## How to Export Gmail into a Mbox File\n",
        "\n",
        "### Label chunks of conversations\n",
        "I choose to export only email threads in which I participated, specifically those in my Sent folder, to remove spam and general announcements. I also choose to export emails in chunks of 6 months to keep the file sizes manageable.\n",
        "\n",
        "1. Filter gmail conversations with, for example, the search \"in:sent after:2019-01-01 before:2019-07-01\". The phrase \"in:sent\" selects only email threads in which I participated by sending a message. The phrase \"after:2019-01-01 before:2019-07-01\" selects only email threads in that date range.\n",
        "\n",
        "2. Check the checkbox to select all conversations in the search result, then click on \"Select all conversations that match this search\".\n",
        "\n",
        "3. Label all these conversations with something like \"in-sent-after-2019-01-01-before-2019-07-01\"\n",
        "\n",
        "### Exporting Gmail (Optionally: only emails with a specific label)\n",
        "\n",
        "4. Go to the Google Export app at [https://takeout.google.com/](https://takeout.google.com/)\n",
        "\n",
        "5. Click \"Deselect all\", then find and check \"Mail: Messages and attachments in your Gmail account in MBOX format.\" to only export emails, not all your Google data.\n",
        "\n",
        "6. To select only one label: click \"All Mail data included\", then deselect \"Include all messages in Mail\", then click \"Select all\", then \"Deselect all\". Now, find and check your label \"in-sent-after-2019-01-01-before-2019-07-01\".\n",
        "\n",
        "7. Find and click \"Next Step\" to choose the file type, frequency, and destination.\n",
        "\n",
        "8. Under \"Transfer to:\", I choose \"Add to Drive\", but you could also \"Send download link via email\" or use another option.\n",
        "\n",
        "9. Click \"Create export\".\n",
        "\n",
        "10. Once the export is created, you can click \"Open in Drive\", then \"Open with ZIP Extractor\". Extraction will create a .MBOX file. Take note where these files are placed in your Drive; they might wind up in \"My Drive/Takeout/Mail\".\n",
        "\n",
        "11. You should now have a MBOX file that can be processed with this Python program."
      ],
      "metadata": {
        "id": "2mz8Xu7d0XwA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Mount to Google Drive and show Drive's mbox files"
      ],
      "metadata": {
        "id": "C28qLU4iSjZZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive', force_remount=True)\n",
        "print()\n",
        "\n",
        "mbox_files = []\n",
        "mbox_json_files = []\n",
        "mbox_txt_files = []\n",
        "\n",
        "import os\n",
        "for x in os.listdir('/content/drive/MyDrive'):\n",
        "    if x.endswith('mbox'):\n",
        "        mbox_files.append(x)\n",
        "    if 'mbox' in x and x.endswith('json'):\n",
        "        mbox_json_files.append(x)\n",
        "    if 'mbox' in x and x.endswith('txt'):\n",
        "        mbox_txt_files.append(x)\n",
        "\n",
        "print(\"Mbox files\")\n",
        "for x in sorted(mbox_files):\n",
        "    print(x)\n",
        "print()\n",
        "\n",
        "print(\"Mbox json files\")\n",
        "for x in sorted(mbox_json_files):\n",
        "    print(x)\n",
        "print()\n",
        "\n",
        "print(\"Mbox txt files\")\n",
        "for x in sorted(mbox_txt_files):\n",
        "    print(x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JJmBR_GWTfM9",
        "outputId": "8f07b3e8-6370-4e27-8348-2bd7bdd5f991"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n",
            "\n",
            "Mbox files\n",
            "greg-egan-two-emails.mbox\n",
            "in-sent-after-2016-01-01-before-2016-07-01.mbox\n",
            "in-sent-after-2016-07-01-before-2016-12-31.mbox\n",
            "in-sent-after-2017-01-01-before-2017-07-01.mbox\n",
            "in-sent-after-2017-07-01-before-2017-12-31.mbox\n",
            "in-sent-after-2018-01-01-before-2018-07-01.mbox\n",
            "in-sent-after-2018-07-01-before-2018-12-31.mbox\n",
            "in-sent-after-2019-01-01-before-2019-07-01.mbox\n",
            "in-sent-after-2019-07-01-before-2019-12-31.mbox\n",
            "in-sent-after-2020-01-01-before-2020-07-01.mbox\n",
            "in-sent-after-2020-07-01-before-2020-12-31.mbox\n",
            "in-sent-after-2021-01-01-before-2021-07-01.mbox\n",
            "in-sent-after-2021-07-01-before-2021-12-31.mbox\n",
            "in-sent-after-2022-01-01-before-2022-07-01.mbox\n",
            "in-sent-after-2022-07-01-before-2022-12-31.mbox\n",
            "in-sent-after-2023-01-01-before-2023-07-01.mbox\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox\n",
            "\n",
            "Mbox json files\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox-2.json\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox-4.json\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox-5.json\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox-raw-1.json\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox-raw-2.json\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox.json\n",
            "\n",
            "Mbox txt files\n",
            "in-sent-after-2016-07-01-before-2016-12-31.mbox-final.txt\n",
            "in-sent-after-2017-01-01-before-2017-07-01.mbox-test-2.txt\n",
            "in-sent-after-2017-07-01-before-2017-12-31.mbox-test-2.txt\n",
            "in-sent-after-2018-01-01-before-2018-07-01.mbox-test-2.txt\n",
            "in-sent-after-2018-07-01-before-2018-12-31.mbox-test-2.txt\n",
            "in-sent-after-2019-01-01-before-2019-07-01.mbox-test-2.txt\n",
            "in-sent-after-2019-07-01-before-2019-12-31.mbox-test-2.txt\n",
            "in-sent-after-2020-01-01-before-2020-07-01.mbox-test-2.txt\n",
            "in-sent-after-2020-07-01-before-2020-12-31.mbox-test-2.txt\n",
            "in-sent-after-2021-01-01-before-2021-07-01.mbox-test-2.txt\n",
            "in-sent-after-2021-07-01-before-2021-12-31.mbox-test-2.txt\n",
            "in-sent-after-2022-01-01-before-2022-07-01.mbox-test-2.txt\n",
            "in-sent-after-2022-07-01-before-2022-12-31.mbox-final-2.txt\n",
            "in-sent-after-2022-07-01-before-2022-12-31.mbox-test-1.txt\n",
            "in-sent-after-2023-01-01-before-2023-07-01.mbox-1.txt\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox-4.txt\n",
            "in-sent-after-2023-07-01-before-2024-01-12.mbox-5.txt\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Check a specific file"
      ],
      "metadata": {
        "id": "cWEEvRWd52EE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "file_name = \"in-sent-after-2016-01-01-before-2016-07-01.mbox\"\n",
        "path = '/content/drive/My Drive'\n",
        "if file_name in os.listdir(path):\n",
        "    print(\"File found\")\n",
        "else:\n",
        "    print(\"File not found\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x3GrQtrigtzz",
        "outputId": "77a3a9c0-c943-4bef-e63b-e1d8ff97d837"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "File found\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Helper functions to read mbox"
      ],
      "metadata": {
        "id": "SA-7F5sWSnY0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def print_current_time():\n",
        "    from datetime import datetime, timedelta\n",
        "    # Format the current date and time in a human-readable format\n",
        "    print((datetime.now()-timedelta(hours=5)).strftime(\"%Y-%m-%d %H:%M:%S\"))\n",
        "\n",
        "def UNIX_timestamp_to_formatted_datetime(date):\n",
        "    unix_timestamp = int(date)  # Convert to integer and then to seconds\n",
        "    date_time_obj = datetime.utcfromtimestamp(unix_timestamp)\n",
        "    # Format the datetime object as a string\n",
        "    formatted_date = date_time_obj.strftime('%Y-%m-%d %H:%M:%S')\n",
        "    return formatted_date\n",
        "\n",
        "def getcharsets(msg):\n",
        "    charsets = set({})\n",
        "    for c in msg.get_charsets():\n",
        "        if c is not None:\n",
        "            charsets.update([c])\n",
        "    return charsets\n",
        "\n",
        "def getBody(msg):\n",
        "    while msg.is_multipart():\n",
        "        msg=msg.get_payload()[0]\n",
        "    t=msg.get_payload(decode=True)\n",
        "    for charset in getcharsets(msg):\n",
        "        t=t.decode(charset)\n",
        "    return t\n",
        "\n",
        "def handleerror(errmsg, emailmsg,cs):\n",
        "    print()\n",
        "    print(errmsg)\n",
        "    print(\"This error occurred while decoding with \",cs,\" charset.\")\n",
        "    print(\"These charsets were found in the one email.\",getcharsets(emailmsg))\n",
        "    print(\"This is the subject:\",emailmsg['subject'])\n",
        "    print(\"This is the sender:\",emailmsg['From'])\n",
        "\n",
        "def getbodyfromemail(msg):\n",
        "    body = None\n",
        "    #Walk through the parts of the email to find the text body.\n",
        "    if msg.is_multipart():\n",
        "        for part in msg.walk():\n",
        "\n",
        "            # If part is multipart, walk through the subparts.\n",
        "            if part.is_multipart():\n",
        "\n",
        "                for subpart in part.walk():\n",
        "                    if subpart.get_content_type() == 'text/plain':\n",
        "                        # Get the subpart payload (i.e the message body)\n",
        "                        body = subpart.get_payload(decode=True)\n",
        "                        #charset = subpart.get_charset()\n",
        "\n",
        "            # Part isn't multipart so get the email body\n",
        "            elif part.get_content_type() == 'text/plain':\n",
        "                body = part.get_payload(decode=True)\n",
        "                #charset = part.get_charset()\n",
        "\n",
        "    # If this isn't a multi-part message then get the payload (i.e the message body)\n",
        "    elif msg.get_content_type() == 'text/plain':\n",
        "        body = msg.get_payload(decode=True)\n",
        "\n",
        "   # No checking done to match the charset with the correct part.\n",
        "    for charset in getcharsets(msg):\n",
        "        try:\n",
        "            body = body.decode(charset)\n",
        "        except UnicodeDecodeError:\n",
        "            continue\n",
        "            # handleerror(\"UnicodeDecodeError: encountered.\",msg,charset)\n",
        "        except AttributeError:\n",
        "            continue\n",
        "            # handleerror(\"AttributeError: encountered\" ,msg,charset)\n",
        "\n",
        "        # body = body.decode('utf-8', errors='replace')\n",
        "\n",
        "    return body\n",
        "\n",
        "from datetime import datetime\n",
        "import time\n",
        "from dateutil import parser\n",
        "\n",
        "def email_datetime_str_to_unix_timestamp(datetime_str):\n",
        "\n",
        "    # The provided datetime string\n",
        "    example_datetime_str = \"Sat, 16 Dec 2023 18:59:31 -0500\"\n",
        "\n",
        "    if not datetime_str:\n",
        "        datetime_str = 1704826771\n",
        "\n",
        "    if type(datetime_str) == int:\n",
        "        return datetime_str\n",
        "\n",
        "    if datetime_str[-6:] in [\" (\"+x+\")\" for x in [\"UTC\", \"PDT\", \"GMT\", \"PST\", \"EST\", \"EDT\", \"MDT\"]]:\n",
        "        datetime_str = datetime_str[:-6]\n",
        "\n",
        "    if datetime_str[-12:] in [\" (\"+x+\")\" for x in [\"GMT+01:00\", \"GMT+08:00\"]]:\n",
        "        datetime_str = datetime_str[:-12]\n",
        "\n",
        "    if datetime_str[3] == ',':\n",
        "        datetime_str = datetime_str[5:]\n",
        "\n",
        "    # Parse the date string into a datetime object\n",
        "    dt = parser.parse(datetime_str)\n",
        "\n",
        "    # Convert to Unix timestamp\n",
        "    unix_timestamp = int(time.mktime(dt.timetuple()))\n",
        "\n",
        "    return unix_timestamp"
      ],
      "metadata": {
        "id": "vFckKDyUtLbT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Designate mbox file name to extract"
      ],
      "metadata": {
        "id": "zsCokQDlSqMY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "mbox_names = [\n",
        "    \"in-sent-after-2016-07-01-before-2016-12-31.mbox\",\n",
        "    \"in-sent-after-2017-01-01-before-2017-07-01.mbox\",\n",
        "    \"in-sent-after-2017-07-01-before-2017-12-31.mbox\",\n",
        "    \"in-sent-after-2018-01-01-before-2018-07-01.mbox\",\n",
        "    \"in-sent-after-2018-07-01-before-2018-12-31.mbox\",\n",
        "    \"in-sent-after-2019-01-01-before-2019-07-01.mbox\",\n",
        "    \"in-sent-after-2019-07-01-before-2019-12-31.mbox\",\n",
        "    \"in-sent-after-2020-01-01-before-2020-07-01.mbox\",\n",
        "    \"in-sent-after-2020-07-01-before-2020-12-31.mbox\",\n",
        "    \"in-sent-after-2021-01-01-before-2021-07-01.mbox\",\n",
        "    \"in-sent-after-2021-07-01-before-2021-12-31.mbox\",\n",
        "    \"in-sent-after-2022-01-01-before-2022-07-01.mbox\",\n",
        "    \"in-sent-after-2022-07-01-before-2022-12-31.mbox\",\n",
        "    \"in-sent-after-2023-01-01-before-2023-07-01.mbox\",\n",
        "    \"in-sent-after-2023-07-01-before-2024-01-12.mbox\",\n",
        "]\n",
        "\n",
        "mbox_name = mbox_names[14]\n",
        "\n",
        "print(f\"mbox_name: {mbox_name}\")\n",
        "\n",
        "# Path to your mbox file\n",
        "mbox_path = '/content/drive/My Drive/'+mbox_name\n",
        "\n",
        "import mailbox\n",
        "\n",
        "# Open the mbox file\n",
        "mbox = mailbox.mbox(mbox_path)"
      ],
      "metadata": {
        "id": "r4jSM7ssulmF",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "319a6a48-00e5-442c-9a27-e6019eb90d4a"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "mbox_name: in-sent-after-2023-07-01-before-2024-01-12.mbox\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Save emails into a simpler data structure"
      ],
      "metadata": {
        "id": "HOERFPxg6HZV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nQIwcQCdYLL0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5ebb6c9e-215e-462c-f1a2-81186306ba30"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of emails: 1932\n",
            "Number of threads: 603\n"
          ]
        }
      ],
      "source": [
        "emails = []\n",
        "thread_ids = []\n",
        "\n",
        "i = 0\n",
        "limit = 10000\n",
        "\n",
        "# Iterate through messages\n",
        "for message in mbox:\n",
        "    # print(message)\n",
        "    # break\n",
        "\n",
        "    try:\n",
        "        # Extracting basic headers\n",
        "        subject = message['subject']\n",
        "        sender = message['from']\n",
        "        receiver = message['to']\n",
        "        date = message['date']\n",
        "        thread_id = message['X-GM-THRID']\n",
        "\n",
        "        # Accessing the body of the email\n",
        "        body_text = getbodyfromemail(message)\n",
        "\n",
        "        if len(body_text) > 4096:\n",
        "            body_text = body_text[:4096]\n",
        "\n",
        "        if body_text[:18] == 'Total SMS messages':\n",
        "            body_text = 'Total SMS messages'\n",
        "\n",
        "        if not body_text:\n",
        "            continue\n",
        "\n",
        "        if not thread_id in thread_ids:\n",
        "            thread_ids.append(thread_id)\n",
        "\n",
        "        json_message = {\n",
        "            'subject': subject,\n",
        "            'sender': sender,\n",
        "            'receiver': receiver,\n",
        "            'date': email_datetime_str_to_unix_timestamp(date),\n",
        "            'message': body_text,\n",
        "            'thread_id': thread_id\n",
        "        }\n",
        "\n",
        "        emails.append(json_message)\n",
        "    except:\n",
        "        continue\n",
        "    i += 1\n",
        "    if i > limit:\n",
        "        break\n",
        "\n",
        "# emails = sorted(emails, key=lambda x: x['date'])\n",
        "\n",
        "print(f\"Number of emails: {len(emails)}\")\n",
        "print(f\"Number of threads: {len(thread_ids)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Check that timestamp decoding is working"
      ],
      "metadata": {
        "id": "uaAcT93S6J6a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for i,x in enumerate(emails):\n",
        "    # Unix timestamp to be converted\n",
        "    timestamp = emails[i]['date']\n",
        "\n",
        "    print(timestamp)\n",
        "\n",
        "    print(UNIX_timestamp_to_formatted_datetime(timestamp))\n",
        "    print()\n",
        "\n",
        "    if i > 2:\n",
        "        break"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "s-PVA0CSVYA5",
        "outputId": "0600a192-c931-483f-8186-42d4a207f9b4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1656852029\n",
            "2022-07-03 12:40:29\n",
            "\n",
            "1656578955\n",
            "2022-06-30 08:49:15\n",
            "\n",
            "1695686942\n",
            "2023-09-26 00:09:02\n",
            "\n",
            "1645786228\n",
            "2022-02-25 10:50:28\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Inspect email data structure"
      ],
      "metadata": {
        "id": "G54R_rE46Noo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for x in emails[:3]:\n",
        "    print(x)"
      ],
      "metadata": {
        "id": "BdwxWHZyZiHb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Clean up sender, receiver, and message"
      ],
      "metadata": {
        "id": "8HbXcFmxa7ZW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "good_emails = []\n",
        "receiver_names = []\n",
        "sender_names = []\n",
        "\n",
        "for e in emails:\n",
        "    try:\n",
        "        subject = e['subject']\n",
        "        sender = e['sender']\n",
        "        receiver = e['receiver']\n",
        "        date = e['date']\n",
        "        message = e['message']\n",
        "        thread_id = e['thread_id']\n",
        "\n",
        "        # print(subject)\n",
        "        # print(sender)\n",
        "        # print(receiver)\n",
        "        # print(date)\n",
        "        # print(message)\n",
        "        # print()\n",
        "\n",
        "        if not receiver:\n",
        "            continue\n",
        "\n",
        "        sender_name = sender.split('<')[0].strip()\n",
        "        sender_names.append(sender_name)\n",
        "\n",
        "        receiver_name = receiver.split('<')[0].strip()\n",
        "\n",
        "        email_receiver_names = []\n",
        "        for n in receiver.split(','):\n",
        "            if '<' in n:\n",
        "                receiver_names.append(receiver_name)\n",
        "                email_receiver_names.append(receiver_name)\n",
        "            else:\n",
        "                receiver_names.append(n)\n",
        "                email_receiver_names.append(n)\n",
        "\n",
        "        if message:\n",
        "            if '\\nOn' in message:\n",
        "                message = message.split('\\nOn')[0].strip()\n",
        "\n",
        "            email_normalized = message.replace('\\r\\n', '\\n')\n",
        "            lines = email_normalized.split('\\n')\n",
        "\n",
        "            # Check if the last line contains the first or last name to remove sender signature\n",
        "            sender_names = sender_name.split(' ')\n",
        "\n",
        "            last_line = lines[-1] if lines else \"\"\n",
        "\n",
        "            while not last_line and lines:\n",
        "                lines = lines[:-1]  # Remove the last line\n",
        "                last_line = lines[-1] if lines else \"\"\n",
        "\n",
        "            if any([x in last_line for x in sender_names]) and lines:\n",
        "                lines = lines[:-1]  # Remove the last line\n",
        "\n",
        "                last_line = lines[-1] if lines else \"\"\n",
        "\n",
        "                # Check if the last line contains one of the following to remove the parting words\n",
        "                parting_words = [\"Best\", \"Thanks\", \"Sincerely\", \"Regards\", \"Warm regards\",\n",
        "                \"Kind regards\", \"Best regards\", \"Respectfully\", \"Cheers\", \"Take care\",\n",
        "                \"All the best\", \"Many thanks\"]\n",
        "\n",
        "                if any([x in last_line for x in parting_words]) and lines:\n",
        "                    lines = lines[:-1]  # Remove the last line\n",
        "\n",
        "            message = ' '.join(lines)\n",
        "\n",
        "            message = message.replace('  ',' ')\n",
        "\n",
        "            json_message = {\n",
        "              'subject': subject,\n",
        "              'sender': sender_name,\n",
        "              'receiver': receiver_name,\n",
        "              'date': date,\n",
        "              'message': message,\n",
        "              'thread_id': thread_id\n",
        "            }\n",
        "\n",
        "            good_emails.append(json_message)\n",
        "\n",
        "    except Exception as x:\n",
        "        errnum = x.args[0]\n",
        "        print(e)\n",
        "        print(errnum)\n",
        "        print()"
      ],
      "metadata": {
        "id": "fFJ7vVcoPcV0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"Number of emails: {len(emails)}\")\n",
        "print(f\"Number of good emails: {len(good_emails)}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FTormRy2Ql1C",
        "outputId": "43a0af87-67ed-4372-f423-aab7e82b96b9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of emails: 1932\n",
            "Number of good emails: 1931\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Group into Threads"
      ],
      "metadata": {
        "id": "BHv458Ikb0_V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "threads = []\n",
        "\n",
        "for t in thread_ids:\n",
        "    thread = []\n",
        "    for e in good_emails:\n",
        "        if e['thread_id'] == t:\n",
        "            thread.append(e)\n",
        "    thread = sorted(thread, key=lambda x: x['date'])\n",
        "    if thread:\n",
        "        threads.append(thread)"
      ],
      "metadata": {
        "id": "91idoRRT4rQb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Check correct grouping of threads"
      ],
      "metadata": {
        "id": "nZ3jU1370RHo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "i = 0\n",
        "limit = 3\n",
        "for t in threads[:10]:\n",
        "    if len(t) < 3:\n",
        "        continue\n",
        "    print(t[0]['subject'])\n",
        "    print(UNIX_timestamp_to_formatted_datetime(t[0]['date']))\n",
        "    for e in t:\n",
        "        print('sender:',e['sender'],',',e['message'])\n",
        "        print()\n",
        "    print()"
      ],
      "metadata": {
        "id": "c3OEVKBJ4FDQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Helper Functions to Write Emails into TXT File"
      ],
      "metadata": {
        "id": "fWnI5CgqSFzb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datetime import datetime\n",
        "\n",
        "def UNIX_timestamp_to_formatted_datetime(date):\n",
        "    unix_timestamp = int(date)  # Convert to integer and then to seconds\n",
        "    date_time_obj = datetime.utcfromtimestamp(unix_timestamp)\n",
        "    # Format the datetime object as a string\n",
        "    formatted_date = date_time_obj.strftime('%Y-%m-%d %H:%M:%S')\n",
        "    return formatted_date\n",
        "\n",
        "def write_threads_to_txt_file(output_file_path, threads):\n",
        "    if not 'txt' in output_file_path:\n",
        "        print(\"txt not in\",output_file_path)\n",
        "        return\n",
        "\n",
        "    # Writing the output to a text file\n",
        "    with open(output_file_path, 'w') as file:\n",
        "        print(\"writing to output_file_path:\", output_file_path)\n",
        "\n",
        "        num_emails = 0\n",
        "        contacts = []\n",
        "        all_thread_contacts = []\n",
        "        earliest = threads[0][0]['date']\n",
        "        latest = threads[0][0]['date']\n",
        "\n",
        "        for t in threads:\n",
        "            thread_contacts = []\n",
        "            for e in t:\n",
        "                contacts.append(e['sender'])\n",
        "                contacts.append(e['receiver'])\n",
        "                thread_contacts.append(e['sender'])\n",
        "                thread_contacts.append(e['receiver'])\n",
        "                num_emails += 1\n",
        "                message_date = e['date']\n",
        "                if e['date'] < earliest:\n",
        "                    earliest = message_date\n",
        "                if e['date'] > latest:\n",
        "                    latest = message_date\n",
        "            all_thread_contacts.append(list(set(thread_contacts)))\n",
        "\n",
        "        contacts = list(set(contacts))\n",
        "\n",
        "        file.write('Total emails: ' + str(num_emails) + '\\n')\n",
        "        file.write('Contacts found: ' + str(len(contacts)) + '\\n')\n",
        "        file.write('Earliest message: ' + UNIX_timestamp_to_formatted_datetime(earliest) + '\\n')\n",
        "        file.write('Latest message: ' + UNIX_timestamp_to_formatted_datetime(latest) + '\\n')\n",
        "\n",
        "        file.write('\\n')\n",
        "\n",
        "        for i,t in enumerate(threads):\n",
        "            thread_contacts = all_thread_contacts[i]\n",
        "            max_length = max(len(x) for x in thread_contacts)\n",
        "\n",
        "            file.write('Conversation with ' + ', '.join(thread_contacts) + '\\n')\n",
        "            padded_contact_names = {x: x.ljust(max_length) for x in thread_contacts}\n",
        "\n",
        "            file.write(str(len(t)) + \" emails\" + '\\n')\n",
        "\n",
        "            date_sorted_emails = sorted(t, key=lambda x: x['date'])\n",
        "\n",
        "            earliest = date_sorted_emails[0]['date']\n",
        "            latest = date_sorted_emails[-1]['date']\n",
        "\n",
        "            file.write('Earliest message: ' + UNIX_timestamp_to_formatted_datetime(earliest) + '\\n')\n",
        "            file.write('Latest message: ' + UNIX_timestamp_to_formatted_datetime(latest) + '\\n')\n",
        "\n",
        "            # Iterate over messages\n",
        "            for e in t:\n",
        "                # Format the datetime object as a string\n",
        "                formatted_date = UNIX_timestamp_to_formatted_datetime(e['date'])\n",
        "                file.write(formatted_date + \": \" + padded_contact_names[e['sender']] + \": \"+e['message'] + '\\n')\n",
        "            file.write('\\n')\n",
        "\n",
        "    print('File closed')"
      ],
      "metadata": {
        "id": "-4Rnc-zlSEet"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Check target mbox name"
      ],
      "metadata": {
        "id": "JWyb4vxw6sDU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "mbox_name"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 36
        },
        "id": "Wux5BAdgBYua",
        "outputId": "b092b3ea-bce5-495a-e4d4-673949ecb0e8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'in-sent-after-2022-01-01-before-2022-07-01.mbox'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 188
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Write Emails into TXT File"
      ],
      "metadata": {
        "id": "i5HjgcNB6vrz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print_current_time()\n",
        "\n",
        "txt_suffix = '-test-2.txt'\n",
        "txt_name = mbox_name + txt_suffix\n",
        "txt_path = mbox_path + txt_suffix\n",
        "\n",
        "print('txt_path:',txt_path)\n",
        "\n",
        "# Print the number of emails\n",
        "print('Total emails:', len(good_emails))\n",
        "\n",
        "if txt_name in os.listdir('/content/drive/My Drive'):\n",
        "    print(\"File already exists\")\n",
        "else:\n",
        "    write_threads_to_txt_file(txt_path,threads)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GsVz-XaqZfI0",
        "outputId": "ed1e6ba5-6af9-4181-d11f-086e7ded0f83"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2024-05-10 19:44:04\n",
            "txt_path: /content/drive/My Drive/in-sent-after-2022-01-01-before-2022-07-01.mbox-test-2.txt\n",
            "Total emails: 1931\n",
            "writing to output_file_path: /content/drive/My Drive/in-sent-after-2022-01-01-before-2022-07-01.mbox-test-2.txt\n",
            "File closed\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Helper Functions to Create Role/System/User/Assistant JSON files"
      ],
      "metadata": {
        "id": "wx7E7XGmSB_M"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def write_threads_to_role_system_user_format_json(output_file_path, threads):\n",
        "    if not 'json' in output_file_path:\n",
        "        print(\"json not in\",output_file_path)\n",
        "        return\n",
        "\n",
        "    # Writing the output to a text file\n",
        "    with open(output_file_path, 'w') as file:\n",
        "        print(\"writing to output_file_path:\", output_file_path)\n",
        "\n",
        "        for t in threads:\n",
        "            if len(t) < 2:\n",
        "                continue\n",
        "\n",
        "            conversation = []\n",
        "\n",
        "            system_message = (\"Be polite and formal. Do not apologize. Use correct grammar and avoid logic fallacies.\")\n",
        "            conversation.append({\"role\": \"system\", \"content\": system_message})\n",
        "\n",
        "            assistant_present = False\n",
        "\n",
        "            # Iterate over messages\n",
        "            for e in t:\n",
        "                sender_name = e['sender']\n",
        "\n",
        "                if \"isaac\" in e['sender'].lower():\n",
        "                    role = \"assistant\"\n",
        "                    assistant_present = True\n",
        "                else:\n",
        "                    role = \"user\"\n",
        "                conversation.append({\"role\": role, \"content\": e['message']})\n",
        "\n",
        "                # conversation.append({\"role\": \"separator\", \"content\": \"<END_OF_CONVERSATION>\"})\n",
        "            if assistant_present:\n",
        "                json_record = json.dumps({'messages': conversation})\n",
        "                file.write(json_record + '\\n')"
      ],
      "metadata": {
        "id": "_cZqT6Ha4FF_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Create Role/System/User/Assistant JSON file"
      ],
      "metadata": {
        "id": "VCuWaoaJ7Ri9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print_current_time()\n",
        "\n",
        "json_suffix = '-5.json'\n",
        "json_file = mbox_name + json_suffix\n",
        "json_path = mbox_path + json_suffix\n",
        "\n",
        "print('json_path:',json_path)\n",
        "\n",
        "# Print the number of emails\n",
        "print('Total emails:', len(good_emails))\n",
        "\n",
        "if json_file in os.listdir('/content/drive/My Drive'):\n",
        "    print(\"File already exists\")\n",
        "else:\n",
        "    write_threads_to_role_system_user_format_json(json_path,threads)"
      ],
      "metadata": {
        "id": "jFohQH-8bt3v",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b94b6806-e153-4170-ffea-707e41012365"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2024-01-15 20:06:43\n",
            "json_path: /content/drive/My Drive/in-sent-after-2023-07-01-before-2024-01-12.mbox-5.json\n",
            "Total emails: 1291\n",
            "writing to output_file_path: /content/drive/My Drive/in-sent-after-2023-07-01-before-2024-01-12.mbox-5.json\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Save threads into json"
      ],
      "metadata": {
        "id": "yBmhVUFEyCLe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print_current_time()\n",
        "\n",
        "import json\n",
        "\n",
        "# File path where JSON data is saved\n",
        "raw_json_suffix = \"-raw-2.json\"\n",
        "\n",
        "raw_json_file = mbox_name + raw_json_suffix\n",
        "raw_json_path = mbox_path + raw_json_suffix\n",
        "\n",
        "print('raw_json_path:',raw_json_path)\n",
        "\n",
        "if raw_json_file in os.listdir('/content/drive/MyDrive'):\n",
        "    print(\"File already exists found\")\n",
        "else:\n",
        "    with open(raw_json_path, 'w') as file:\n",
        "        json.dump(threads, file)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NuFazGmakUgL",
        "outputId": "1a061571-baee-4e2a-b8e0-5ace07b71aad"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2024-01-15 20:00:16\n",
            "raw_json_path: /content/drive/My Drive/in-sent-after-2023-07-01-before-2024-01-12.mbox-raw-2.json\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load threads from json"
      ],
      "metadata": {
        "id": "kZHIG26Y86ov"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "raw_json_file = mbox_name + \"-raw-1.json\"\n",
        "raw_json_path = mbox_path + \"-raw-1.json\"\n",
        "\n",
        "print('raw_json_file:',raw_json_file)\n",
        "\n",
        "if not raw_json_file in os.listdir('/content/drive/My Drive'):\n",
        "    print(\"File doesn't exist\")\n",
        "else:\n",
        "    print(\"File found\")\n",
        "    with open(raw_json_path, 'r') as file:\n",
        "        threads = json.load(file)\n",
        "\n",
        "print(f\"Number of threads: {len(threads)}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "knJ-HfTV6vw1",
        "outputId": "4db496be-14b6-4214-b3ee-97d723f3f04f"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "raw_json_file: in-sent-after-2023-07-01-before-2024-01-12.mbox-raw-1.json\n",
            "File found\n",
            "Number of threads: 79\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}