{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bag-of-words, Tf-idf, and cosine similarity\n",
    "\n",
    "In this notebook, we will use the bag-of-words model, the tf-idf model, and the cosine similarity to compare the similarity between the plays written by William Shakespeare. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import os\n",
    "\n",
    "# Load the text file\n",
    "with open('data/shakespeare.txt', 'r', encoding='utf-8') as file:\n",
    "    text = file.read()\n",
    "\n",
    "# Define the regular expression pattern\n",
    "play_start_pattern = re.compile(\n",
    "    r'^\\s*(\\d{4})\\s*'                        # Year line\n",
    "    r'\\n\\s*([A-Z][A-Z\\s\\'.,;:!?\\-\\&]+)\\s*'   # Title line\n",
    "    r'(?:\\n\\s*by William Shakespeare\\s*)?',   # Optional author line\n",
    "    flags=re.MULTILINE\n",
    ")\n",
    "\n",
    "# Find all matches\n",
    "matches = list(play_start_pattern.finditer(text))\n",
    "\n",
    "# Split the text into plays\n",
    "play_dict = {}\n",
    "for i, match in enumerate(matches):\n",
    "    year = match.group(1).strip()\n",
    "    title = match.group(2).strip()\n",
    "    start_index = match.end()\n",
    "    end_index = matches[i + 1].start() if i + 1 < len(matches) else len(text)\n",
    "    content = text[start_index:end_index].strip()\n",
    "    full_title = f\"{year} - {title}\"\n",
    "    play_dict[full_title] = content\n",
    "\n",
    "# Print the plays detected\n",
    "print(f\"\\n{len(play_dict)} plays detected:\")\n",
    "for title in play_dict.keys():\n",
    "    print(title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert play_dict into a DataFrame\n",
    "plays_df = pd.DataFrame(list(play_dict.items()), columns=['Title', 'Content'])\n",
    "\n",
    "plays_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocess the Text Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "from nltk.corpus import stopwords\n",
    "import string\n",
    "\n",
    "# Download stopwords if you haven't already\n",
    "#nltk.download('stopwords')\n",
    "\n",
    "# Define a function to preprocess text\n",
    "def preprocess_text(text):\n",
    "    # Lowercase the text\n",
    "    text = text.lower()\n",
    "    # Remove punctuation\n",
    "    text = text.translate(str.maketrans('', '', string.punctuation))\n",
    "    # Tokenize the text\n",
    "    tokens = text.split()\n",
    "    # Remove stop words\n",
    "    stop_words = set(stopwords.words('english'))\n",
    "    tokens = [word for word in tokens if word not in stop_words]\n",
    "    # Remove numbers\n",
    "    tokens = [word for word in tokens if not word.isnumeric()]\n",
    "    # Rejoin tokens into a single string\n",
    "    return ' '.join(tokens)\n",
    "\n",
    "# Apply the preprocessing function to the 'Content' column\n",
    "plays_df['Content_clean'] = plays_df['Content'].apply(preprocess_text)\n",
    "plays_df.head()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The document term matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "\n",
    "# Create an instance of CountVectorizer\n",
    "vectorizer = CountVectorizer()\n",
    "# Can also iclude stop word removal directly in the vectorizer\n",
    "#vectorizer = CountVectorizer(stop_words='english')\n",
    "\n",
    "# Fit the vectorizer on the content and transform the content into a document-term matrix\n",
    "dtm = vectorizer.fit_transform(plays_df['Content_clean'])\n",
    "\n",
    "# Get the feature names (i.e., the vocabulary)\n",
    "feature_names = vectorizer.get_feature_names_out()\n",
    "\n",
    "# Convert the document-term matrix to a DataFrame for easier handling\n",
    "dtm_df = pd.DataFrame(dtm.toarray(), columns=feature_names, index=plays_df['Title'])\n",
    "\n",
    "dtm_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtm_df[['william', 'hamlet']].sample(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute the TF-IDF matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "\n",
    "# Create an instance of TfidfVectorizer\n",
    "tfidf_vectorizer = TfidfVectorizer()\n",
    "\n",
    "# Fit and transform the content\n",
    "tfidf = tfidf_vectorizer.fit_transform(plays_df['Content_clean'])\n",
    "\n",
    "# Get feature names\n",
    "tfidf_feature_names = tfidf_vectorizer.get_feature_names_out()\n",
    "\n",
    "# Convert to DataFrame\n",
    "tfidf_df = pd.DataFrame(tfidf.toarray(), columns=tfidf_feature_names, index=plays_df['Title'])\n",
    "tfidf_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sum TF-IDF scores for each term across all documents\n",
    "aggregate_tfidf = tfidf_df.sum(axis=0)\n",
    "aggregate_tfidf = aggregate_tfidf.sort_values(ascending=False)\n",
    "aggregate_tfidf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aggregate_tfidf.index = range(len(aggregate_tfidf))\n",
    "\n",
    "# Plot the aggregate TF-IDF scores\n",
    "plt.figure(figsize=(6, 4))\n",
    "np.log(aggregate_tfidf).plot(lw=2)\n",
    "plt.title('Aggregate TF-IDF Scores')   \n",
    "plt.xlabel('TF-IDF Score')  \n",
    "plt.ylabel('Term')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating Cosine Similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "# Compute cosine similarity matrix\n",
    "cosine_sim_matrix = cosine_similarity(tfidf_df)\n",
    "\n",
    "# Convert to DataFrame for better readability\n",
    "cosine_sim_df = pd.DataFrame(cosine_sim_matrix, index=plays_df['Title'], columns=plays_df['Title'])\n",
    "\n",
    "# Display similarity between plays\n",
    "cosine_sim_df[['1604 - THE TRAGEDY OF HAMLET, PRINCE OF DENMARK']].sample(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (nlp)",
   "language": "python",
   "name": "nlp"
  },
  "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}