# %% [markdown]
"""
# cleaning
big_data.csv.zst is the main file we will be using. Every step in the pipeline adds a new column and overwrites the file. This is reversible and when any step changes everything can be run again regardless og the state of the file.
"""

# %%
import nltk
import re
import os
import time
import pandas as pd

DATA_DIR = "../data"

# %%
# download nltk data
nltk.download("all")

# %%
news_sample = pd.read_csv(f"{DATA_DIR}/news_sample.csv")

# %%
# We will not waste space on csv files, L.
if (os.path.exists(f"{DATA_DIR}/995,000_rows.csv")):
    big_data = pd.read_csv(f"{DATA_DIR}/995,000_rows.csv", low_memory=False)
    big_data.to_csv(f"{DATA_DIR}/big_data.csv.zst")
    os.remove(f"{DATA_DIR}/995,000_rows.csv")
big_data = None

# %%
# cleans text and returns a list of tokens.
def clean_text(
        text,
        remove_regex_patterns = True,
        remove_stopwords = True,
        remove_special_characters = True,
        stemming = True):

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

    if remove_regex_patterns:
        url_pattern = r'\S+\.\S+'
        email_pattern = r'\w+@\w+\.\w+'
        date_pattern = r'[a-z]+ \d{1,2}[a-z]?, \d{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)

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

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

    return tokens

# %% [markdown]
"""
## Output
Now we check what the function does and how the vocabulary changes.
"""

# %%
# Generates a vocabulary (set of unique words) from a pandas series.
def generate_vocabulary(series):
    vocabulary = set()
    series.apply(lambda tokens: vocabulary.update(tokens))
    return vocabulary

# %%
print("original text:\n")
print(news_sample['content'][1])
print("\n" + "-" * 100 + "\n")

print("cleaned tokens:\n")
print(clean_text(news_sample['content'][1]))
print("\n" + "-" * 100 + "\n")

tokenization_size = len(generate_vocabulary(news_sample['content'].apply(clean_text, remove_stopwords = False, stemming = False)))
stopwords_size = len(generate_vocabulary(news_sample['content'].apply(clean_text, remove_stopwords = True, stemming = False)))
stemming_size = len(generate_vocabulary(news_sample['content'].apply(clean_text, remove_stopwords = True, stemming = True)))

print("Unique words after tokenization:")
print(tokenization_size)
print("\nUnique words after stopword removal:")
print(stopwords_size)
print("\nUnique words after stemming:")
print(stemming_size)
print("\nStemming reduction rate:")
print(f"{round(1 - stemming_size / stopwords_size, 4) * 100}%")

# %% [markdown]
"""
## Big Data
Now we clean the big dataset and save it to csv.zst file. Pandas can save and load zstd files just fine, and since it's realtime compression it doesn't really take more time while heavily reducing the file size.
"""

# %%
start = time.perf_counter()
first = True
for big_data in pd.read_csv(f"{DATA_DIR}/big_data.csv.zst", chunksize=10000):
    big_data['tokens'] = big_data['content'].apply(clean_text)
    if first:
        big_data.to_csv(f"{DATA_DIR}/big_data_new.csv.zst", mode='w')
        first = False
    else:
        big_data.to_csv(f"{DATA_DIR}/big_data_new.csv.zst", mode='a')
os.rename(f"{DATA_DIR}/big_data_new.csv.zst",f"{DATA_DIR}/big_data.csv.zst")
print(f"cleaning took {round((time.perf_counter() - start) / 60, 5)} minutes")
