Code_search.py

import sqlite3
import string
import math
from pathlib import Path
from collections import Counter
from porterstemmer import PorterStemmer


DB_PATH = Path("webcrawler.db")
stemmer = PorterStemmer()
# Using the same stop-word list from Code_indexer.py
STOP_WORDS = {
    "a", "about", "after", "again", "all", "am", "an", "and", "any", "are",
    "as", "at", "be", "because", "been", "before", "being", "between",
    "both", "but", "by", "can", "did", "do", "does", "doing", "down",
    "during", "each", "few", "for", "from", "had", "has", "have", "having",
    "he", "her", "here", "him", "his", "how", "i", "if", "in", "into",
    "is", "it", "its", "itself", "just", "me", "more", "most", "my",
    "no", "nor", "not", "now", "of", "off", "on", "once", "only", "or",
    "other", "our", "out", "over", "own", "same", "she", "should", "so",
    "some", "such", "than", "that", "the", "their", "them", "then",
    "there", "these", "they", "this", "those", "through", "to", "too",
    "under", "until", "up", "very", "was", "we", "were", "what", "when",
    "where", "which", "while", "who", "why", "with", "would", "you",
    "your"
}


def connect_database(db_path):
    # Connect to the SQLite database created in Part 2.
    if not db_path.exists():
        raise FileNotFoundError(f"Database file not found: {db_path.resolve()}")
    return sqlite3.connect(db_path)


def clean_query_token(token):
    # Clean one query token using the preprocessing rules from the indexer
    token = token.lower().strip()

    if token == "":
        return None
    # Reject tokens that begin with punctuation
    if token[0] in string.punctuation:
        return None
    # Strip trailing punctuation
    token = token.rstrip(string.punctuation)
    if token == "":
        return None
    # Reject numeric tokens
    if token.isnumeric():
        return None
    # Reject stop words
    if token in STOP_WORDS:
        return None
    # Reject short terms
    if len(token) <= 2:
        return None
    # Apply porter stemming here
    token = stemmer.stem(token, 0, len(token) - 1)

    return token


def process_query(query):
    # Process a bag of words query into cleaned & stemmed terms
    processed_terms = []

    for raw_token in query.split():
        cleaned_token = clean_query_token(raw_token)
        if cleaned_token is not None:
            processed_terms.append(cleaned_token)
    return processed_terms


def lookup_query_terms(cur, processed_terms):
    # Look up each unique term in TermDictionary
    # Returns:
    #   query_terms_freqs: Counter of processed terms
    #   query_terms: dictionary mapping term text to TermId
    #   missing_terms: set of processed terms not found in TermDictionary

    query_term_freqs = Counter(processed_terms)
    query_terms = {}
    missing_terms = []

    for term in query_term_freqs:
        cur.execute(
            "SELECT TermId FROM TermDictionary WHERE Term = ?",
            (term,)
        )

        row = cur.fetchone()

        if row is None:
            missing_terms.append(term)
        else:
            query_terms[term] = row[0]

    return query_term_freqs, query_terms, missing_terms        


def print_no_results(original_query, processed_terms, query_terms, missing_terms):
    # Provide output when no documents or terms can be returned for a query
    print("\nSearch Summary")
    print("==============")
    print(f"Search terms entered: {original_query}")
    print(f"Processed query terms: {processed_terms}")

    print("\nQuery term details:")
    for term in processed_terms:
        if term in query_terms:
            print(f"{term}: found in TermDictionary as TermId {query_terms[term]}")
        else:
            print(f"{term}: not found in TermDictionary")

    print("\nTotal candidate documents retrieved: 0")

    print("\nTop documents:")
    print("No documents contain all query terms.")

    print("\nSimpson algorithm / cosine similarity output:")
    print("No similarity scores calculated as no candidate documents were retrieved.")


def retrieve_postings(cur, query_terms):
    """
    Retrieve postings for each query term.

    Returns:
        postings_by_term: dictionary where each term maps to another dictionary:
            term -> {
                doc_id -> {
                    "tfidf": value,
                    "docfreq": value,
                    "termfreq": value,
                    "termid": value
                }
            }
    """
    postings_by_term = {}

    for term, term_id in query_terms.items():
        cur.execute(
            """
            SELECT DocId, tfidf, docfreq, termfreq
            FROM Posting
            WHERE TermId = ?
            ORDER BY DocId
            """,
            (term_id,)
        )

        postings_by_term[term] = {}
        for row in cur.fetchall():
            doc_id = row[0]
            postings_by_term[term][doc_id] = {
                "tfidf": row[1],
                "docfreq": row[2],
                "termfreq": row[3],
                "termid": term_id
            }
    return postings_by_term


def find_candidate_documents(postings_by_term):
    """
    Find documents that contain all query terms.
    
    Returns:
        A sorted list of candidate DocIds.
    """

    doc_sets = []
    for term, postings in postings_by_term.items():
        doc_sets.append(set(postings.keys()))

    if len(doc_sets) == 0:
        return []
    
    candidate_docs = set.intersection(*doc_sets)
    return sorted(candidate_docs)

def get_total_documents(cur):
    """
    Return total number of documents in the collection,
    """
    cur.execute("SELECT COUNT(*) FROM DocumentDictionary")
    return cur.fetchone()[0]


def calculate_query_weights(cur, query_term_freqs, query_terms, total_documents):
    """
    Calculate tf-idf weights for the query terms.
    Query weight formula:
        query_tfidf = query_term_frequency * log(total_documents / document_frequency)
        
    Returns:
        query weights: dictionary mapping TermId to query tf-idf weight
        query_length: Euclidean length of the query vector
        """
    query_weights = {}

    for term, term_id in query_terms.items():
        cur.execute(
            """
            SELECT docfreq
            FROM Posting
            WHERE TermId = ?
            LIMIT 1
            """,
            (term_id,)
        )

        row = cur.fetchone()
        if row is None:
            continue

        doc_freq = row[0]
        idf = math.log(total_documents / doc_freq)
        query_tf = query_term_freqs[term]
        query_tfidf = query_tf * idf
        query_weights[term_id] = query_tfidf

    query_length = math.sqrt(
        sum(weight ** 2 for weight in query_weights.values())
    )
    return query_weights, query_length


def calculate_document_length(cur, doc_id):
    """
    Calculate the Euclidean length of a document vector.
    
    This uses all tf-idf weights stored for the document in the Posting table,
    Not only query terms.
    """
    cur.execute(
        """
        SELECT tfidf
        FROM Posting
        WHERE DocId = ?
        """,
        (doc_id,)
    )

    squared_sum = 0.0

    for row in cur.fetchall():
        tfidf = row[0]
        squared_sum += tfidf ** 2

    return math.sqrt(squared_sum)


def calculate_cosine_scores(cur, candidate_docs, postings_by_term, query_weights, query_length):
    """
    Calculate cosine similarity scores for each candidate document.
    
    Cosine similarity:
    dot_product(query, document) / (query_length * document_length)
    """
    scores = []

    for doc_id in candidate_docs:
        dot_product = 0.0

        for term, postings in postings_by_term.items():
            if doc_id in postings:
                term_id = postings[doc_id]["termid"]
                document_tfidf = postings[doc_id]["tfidf"]
                query_tfidf = query_weights.get(term_id, 0.0)

                dot_product += query_tfidf * document_tfidf

        document_length = calculate_document_length(cur, doc_id)

        if query_length == 0 or document_length == 0:
            cosine_score = 0.0
        else:
            cosine_score = dot_product / (query_length * document_length)

        scores.append((doc_id, cosine_score))

    scores.sort(key=lambda item: item[1], reverse=True)

    return scores


def get_document_name(cur, doc_id):
    """
    Retrieve the document file name for a given DocId.
    """
    cur.execute(
        """
        SELECT DocumentName
        FROM DocumentDictionary
        WHERE DocId = ?
        """,
        (doc_id,)
    )

    row = cur.fetchone()
    if row is None:
        return f"Unknown document for DocId {doc_id}"
    return row[0]


def print_ranked_results(cur, cosine_scores, total_candidates):
    """
    Print the top 20 ranked documents
    """
    print("\nTop documents:")
    print("==============")

    if len(cosine_scores) == 0:
        print("No documents contain all query terms.")
        return
    
    for rank, item in enumerate(cosine_scores[:20], start=1):
        doc_id, score = item
        document_name = get_document_name(cur, doc_id)
        
        print(f"\nRank: {rank}")
        print(f"Document file name: {document_name}")
        print(f"Cosine similarity score: {score:.6f}")
        print(f"Total candidate documents retrieved: {total_candidates}")



def main():
    print("Using database:", DB_PATH.resolve())
    
    con = connect_database(DB_PATH)
    cur = con.cursor()

    original_query = input("Enter search terms separated by spaces: ")
    processed_terms = process_query(original_query)

    if len(processed_terms) == 0:
        print("\nSearch Summary")
        print("==============")
        print(f"Search terms entered: {original_query}")
        print("Processed query terms: []")
        print("\nTotal candidate documents retrieved: 0")
        print("\nTop documents:")
        print("No valid query terms were entered.")
        con.close()
        return

    query_term_freqs, query_terms, missing_terms = lookup_query_terms(
        cur,
        processed_terms
    )

    if missing_terms:
        print_no_results(
            original_query,
            processed_terms,
            query_terms,
            missing_terms
        )
        con.close()
        return

    print("\nSearch Summary")
    print("==============")
    print(f"Search terms entered: {original_query}")
    print(f"Processed query terms: {processed_terms}")

    print("\nQuery term details:")
    for term, term_id in query_terms.items():
        print(f"{term}: found in TermDictionary as TermId {term_id}")

    postings_by_term = retrieve_postings(cur, query_terms)
    candidate_docs = find_candidate_documents(postings_by_term)

    print(f"\nTotal candidate documents retrieved: {len(candidate_docs)}")

    if len(candidate_docs) == 0:
        print("\nTop documents:")
        print("No documents contain all query terms.")

        print("\nSimpson algorithm / cosine similarity output:")
        print("No similarity scores calculated as no candidate documents were retrieved.")

        con.close()
        return

    total_documents = get_total_documents(cur)

    query_weights, query_length = calculate_query_weights(
        cur,
        query_term_freqs,
        query_terms,
        total_documents
    )

    cosine_scores = calculate_cosine_scores(
        cur,
        candidate_docs,
        postings_by_term,
        query_weights,
        query_length
    )

    print("\nSimpson algorithm / cosine similarity output:")
    print("Ranked results are sorted from highest cosine similarity to lowest.")

    print_ranked_results(
        cur,
        cosine_scores,
        len(candidate_docs)
    )

    con.close()


if __name__ == "__main__":
    main()