from constants import DATASET_DIR, TEMP_DIR
from helper import dataset_iterator
import pyarrow as pa
import pyarrow.parquet as pq
import nltk
import re
import shutil

# cleans text and returns a list of tokens.
def clean_text(
        text:str,
        remove_regex_patterns:bool = True,
        remove_stopwords:bool = True,
        remove_special_characters:bool = True,
        stemming:bool = True
    ) -> list[str]:

    text = str(text).lower().strip()

    if remove_regex_patterns:
        url_pattern = r'https?://\S+|www\.\S+'
        email_pattern = r'[\w.-]+@[\w]+\.[\w]+'
        date_pattern = r'([a-z]+ \d{1,2}[a-z]?, \d{4}|\d{2,4}[-/]\d{2,4}[-/]\d{2,4})'  # add more date patterns
        number_pattern = r'\d+'

        text = re.sub(url_pattern, "<URL>", text)
        text = re.sub(email_pattern, "<EMAIL>", text)
        text = re.sub(date_pattern, "<DATE>", text)
        text = re.sub(number_pattern, "<NUMBER>", text)

    if remove_special_characters:
        text = re.sub(r'[^\w (?:<\w+>)]', " ", text)

    tokenizer = nltk.RegexpTokenizer(r'<\w+>|\w+')
    tokens = tokenizer.tokenize(text) # type: ignore

    if remove_stopwords:
        stopwords = stopwords = nltk.corpus.stopwords.words('english')
        tokens = [token for token in tokens if token not in stopwords] # type: ignore

    if stemming:
        stemmer = nltk.SnowballStemmer("english")
        tokens = [stemmer.stem(token) if not re.match(r'<\w+>', token) else token for token in tokens] # type: ignore

    return tokens # type: ignore

def clean_dataset(filename:str) -> None:
    output_path = f"{TEMP_DIR}/{filename}"
    writer = None
    for chunk in dataset_iterator(f"{DATASET_DIR}/{filename}"):
        chunk['tokens'] = chunk['content'].apply(clean_text)

        columns_in_chunk = chunk.columns #[c for c in chunk.columns]
        table = pa.Table.from_pandas(chunk[columns_in_chunk])
        if writer is None:
            writer = pq.ParquetWriter(output_path, table.schema)
        writer.write_table(table)
    writer.close()

    shutil.move(output_path, f"{DATASET_DIR}/{filename}")

def compute_vocab_reduction(filename: str) -> dict[str, float | int]:

    dataset_path = f"{DATASET_DIR}/{filename}"

    vocab_before_stopwords: set[str] = set()
    vocab_after_stopwords: set[str] = set()

    vocab_after_stemming: set[str] = set()

    for chunk in dataset_iterator(dataset_path):
        contents = chunk["content"]
        for text in contents:
            vocab_before_stopwords.update(clean_text(text, remove_stopwords=False, stemming=False))
            vocab_after_stopwords.update(clean_text(text, remove_stopwords=True, stemming=False))
            vocab_after_stemming.update(clean_text(text, remove_stopwords=True, stemming=True))

    before_stop_size = len(vocab_before_stopwords)
    after_stop_size = len(vocab_after_stopwords)
    before_stem_size = after_stop_size  
    after_stem_size = len(vocab_after_stemming)

    stopwords_reduction_rate = (
        (before_stop_size - after_stop_size) / before_stop_size if before_stop_size else 0.0
    )
    stemming_reduction_rate = (
        (before_stem_size - after_stem_size) / before_stem_size if before_stem_size else 0.0
    )

    return {
        "vocab_size_before_stopwords": before_stop_size,
        "vocab_size_after_stopwords": after_stop_size,
        "stopwords_reduction_rate": stopwords_reduction_rate,
        "vocab_size_before_stemming": before_stem_size,
        "vocab_size_after_stemming": after_stem_size,
        "stemming_reduction_rate": stemming_reduction_rate,
    }