import os
import json
import nltk
import pickle
import string
from nltk.corpus import stopwords
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline

# Download NLTK stopwords
nltk.download("stopwords")

# Path Configuration
main_dir = ""
context_dir = os.path.join(main_dir, "context")
saved_model_dir = os.path.join(main_dir, "saved_model")
tfidf_dir = os.path.join(main_dir, "tfidf")

# Load TF-IDF matrix and vectorizer
tfidf_matrix_filepath = os.path.join(tfidf_dir, "tfidf_matrix.pkl")
tfidf_vectorizer_filepath = os.path.join(tfidf_dir, "tfidf_vectorizer.pkl")
tfidf_matrix = pickle.load(open(tfidf_matrix_filepath, "rb"))
tfidf_vectorizer = pickle.load(open(tfidf_vectorizer_filepath, "rb"))


# TF-IDF Preprocessing
def tfidf_preprocess(text):
    punctuation = string.punctuation
    stop_words = set(stopwords.words("english"))
    text = text.lower()
    text = text.translate(str.maketrans("", "", punctuation))
    text = " ".join([word for word in text.split() if word not in stop_words])
    return text


# Load document texts
documents = [open(os.path.join(context_dir, f)).read() for f in os.listdir(context_dir)]
contexts = [
    line for document in documents for line in document.split("\n") if line != ""
]

# Get document names with line numbers
document_context = []
for f in os.listdir(context_dir):
    lines = open(os.path.join(context_dir, f)).read().split("\n")
    for i, line in enumerate(lines):
        if line != "":
            document_context.append(f"{f[:-4]}_{i}{f[-4:]}")

# Model Configuration
MODEL_NAME = [
    "saved_model/distilbert-base-uncased-distilled-squad_5e-06_16",
    "saved_model/roberta-base-squad2_5e-06_16",
]