"""
NLP Text Lab - Gradio app for Hugging Face Spaces
Tasks: Tokenization, Stop-word removal, Stemming, Lemmatization,
       POS tagging + word frequency, Full preprocessing pipeline.
"""
# NOTE: `spaces` must be imported FIRST (before spacy/torch), otherwise ZeroGPU
# raises "CUDA has been initialized before importing the spaces package".
try:
    import spaces
except ImportError:  # running locally / on CPU hardware
    spaces = None

import string
from collections import Counter

import gradio as gr
import nltk
import pandas as pd
import spacy
from nltk.corpus import stopwords, wordnet
from nltk.stem import (LancasterStemmer, PorterStemmer, SnowballStemmer,
                       WordNetLemmatizer)
from nltk.tokenize import sent_tokenize, word_tokenize

# ------------------------------------------------- ZeroGPU compatibility ----
# ZeroGPU needs at least one @spaces.GPU function. The app does not use a GPU,
# so this is only a harmless placeholder.
if spaces is not None:
    @spaces.GPU
    def _gpu_placeholder():
        return "ok"

# ---------------------------------------------------------------- setup ----
for pkg in ["punkt", "punkt_tab", "stopwords", "wordnet", "omw-1.4",
            "averaged_perceptron_tagger", "averaged_perceptron_tagger_eng"]:
    nltk.download(pkg, quiet=True)


def load_spacy():
    try:
        return spacy.load("en_core_web_sm")
    except OSError:
        from spacy.cli import download
        download("en_core_web_sm")
        return spacy.load("en_core_web_sm")


nlp = load_spacy()
STOP_WORDS = set(stopwords.words("english"))
porter, lancaster = PorterStemmer(), LancasterStemmer()
snowball = SnowballStemmer("english")
wnl = WordNetLemmatizer()

DEFAULT_TEXT = (
    "Consistency is the key to success.\n"
    "Every small step we take today brings us closer to the future we dream of.\n"
    "The journey may be difficult, and sometimes we may feel tired or lost,\n"
    "but we must keep moving forward. Growth takes time, patience, and courage."
)


# -------------------------------------------------------------- helpers ----
def _clean(text):
    return (text or "").strip()


def _is_punct(tok):
    return all(ch in string.punctuation for ch in tok)


def _wn_pos(tag):
    """Convert Penn Treebank tag to WordNet POS."""
    if tag.startswith("J"):
        return wordnet.ADJ
    if tag.startswith("V"):
        return wordnet.VERB
    if tag.startswith("R"):
        return wordnet.ADV
    return wordnet.NOUN


# ---------------------------------------------------------------- tasks ----
def task_tokenize(text):
    text = _clean(text)
    if not text:
        return "Please enter a sentence.", "", "", ""
    words = word_tokenize(text)
    sents = sent_tokenize(text)
    doc = nlp(text)
    spacy_tokens = [t.text for t in doc if not t.is_space]
    spacy_sents = [s.text.strip() for s in doc.sents]

    nltk_out = (f"WORD TOKENS ({len(words)}):\n{words}\n\n"
                f"SENTENCES ({len(sents)}):\n" +
                "\n".join(f"{i+1}. {s}" for i, s in enumerate(sents)))
    spacy_out = (f"WORD TOKENS ({len(spacy_tokens)}):\n{spacy_tokens}\n\n"
                 f"SENTENCES ({len(spacy_sents)}):\n" +
                 "\n".join(f"{i+1}. {s}" for i, s in enumerate(spacy_sents)))
    summary = (f"No. of tokens (NLTK): {len(words)}  |  "
               f"No. of sentences (NLTK): {len(sents)}")
    return summary, nltk_out, spacy_out, ""


def task_stopwords(text, remove_punct, lowercase):
    text = _clean(text)
    if not text:
        return "Please enter a sentence.", "", "", ""
    tokens = word_tokenize(text)
    kept, removed = [], []
    for tok in tokens:
        if remove_punct and _is_punct(tok):
            removed.append(tok)
            continue
        if tok.lower() in STOP_WORDS:   # case-insensitive match
            removed.append(tok)
        else:
            kept.append(tok.lower() if lowercase else tok)
    summary = (f"Original tokens: {len(tokens)}  |  Removed: {len(removed)}  |  "
               f"Remaining: {len(kept)}")
    return summary, str(tokens), str(removed), " ".join(kept)


def task_stemming(text):
    text = _clean(text)
    if not text:
        return pd.DataFrame(columns=["Word", "Porter", "Lancaster", "Snowball"])
    toks = [t for t in word_tokenize(text) if not _is_punct(t)]
    rows = [[t, porter.stem(t), lancaster.stem(t), snowball.stem(t)] for t in toks]
    return pd.DataFrame(rows, columns=["Word", "Porter", "Lancaster", "Snowball"])


def task_lemmatize(text):
    text = _clean(text)
    empty1 = pd.DataFrame(columns=["Word", "POS tag", "Lemma (default=noun)",
                                   "Lemma (POS-aware)"])
    empty2 = pd.DataFrame(columns=["Word", "POS", "spaCy Lemma"])
    if not text:
        return empty1, empty2
    toks = [t for t in word_tokenize(text) if not _is_punct(t)]
    rows = []
    for tok, tag in nltk.pos_tag(toks):
        rows.append([tok, tag, wnl.lemmatize(tok.lower()),
                     wnl.lemmatize(tok.lower(), _wn_pos(tag))])
    doc = nlp(text)
    rows2 = [[t.text, t.pos_, t.lemma_] for t in doc
             if not t.is_space and not t.is_punct]
    return (pd.DataFrame(rows, columns=empty1.columns),
            pd.DataFrame(rows2, columns=empty2.columns))


def task_pos_freq(text):
    text = _clean(text)
    if not text:
        return (pd.DataFrame(columns=["Word", "POS tag"]),
                pd.DataFrame(columns=["Word", "Count"]))
    toks = [t for t in word_tokenize(text) if not _is_punct(t)]
    pos_df = pd.DataFrame(nltk.pos_tag(toks), columns=["Word", "POS tag"])
    words = [t.lower() for t in toks if t.lower() not in STOP_WORDS]
    freq_df = pd.DataFrame(Counter(words).most_common(),
                           columns=["Word", "Count"])
    return pos_df, freq_df


def task_pipeline(text):
    text = _clean(text)
    if not text:
        return ["Please enter a sentence."] + [""] * 5
    tokens = word_tokenize(text.lower())
    no_punct = [t for t in tokens if not _is_punct(t)]
    no_stop = [t for t in no_punct if t not in STOP_WORDS]
    stems = [porter.stem(t) for t in no_stop]
    lemmas = [wnl.lemmatize(t, _wn_pos(tag)) for t, tag in nltk.pos_tag(no_stop)]
    return (f"1. Lowercase + Tokenize\n{tokens}", f"2. Remove punctuation\n{no_punct}",
            f"3. Remove stop words\n{no_stop}", f"4. Stemming (Porter)\n{stems}",
            f"5. Lemmatization (POS-aware)\n{lemmas}",
            " ".join(lemmas))


LEARN_MD = """
## Stemming ("trimming") vs Lemmatization

**Stemming** chops the ending (suffix) off a word using simple rules, so the
result may *not* be a real word. It is fast but rough.

| Word | Stem (Porter) |
|---|---|
| studies | studi |
| running | run |
| difficult | difficult |
| happily | happili |

**Lemmatization** uses a dictionary (WordNet) and the word's part of speech to
return the real base form (the *lemma*). Slower but accurate.

| Word | POS | Lemma |
|---|---|---|
| studies | noun/verb | study |
| running | verb | run |
| better | adjective | good |
| mice | noun | mouse |

**When to use what?** Stemming for search engines/speed, lemmatization when
meaning matters (chatbots, sentiment analysis, NLP models).

## Stop words
Very common words (the, is, to, we ...) that carry little meaning. They are
removed to keep only the useful words.
"""

# ------------------------------------------------------------------- UI ----
with gr.Blocks(title="NLP Text Lab") as demo:
    gr.Markdown("# 🧪 NLP Text Lab\nType a sentence/paragraph, then open any "
                "tab and press the button.")
    text_in = gr.Textbox(label="Your text", lines=5, value=DEFAULT_TEXT)

    with gr.Tab("1. Tokenization"):
        b1 = gr.Button("Tokenize", variant="primary")
        s1 = gr.Textbox(label="Summary")
        with gr.Row():
            n1 = gr.Textbox(label="NLTK", lines=12)
            sp1 = gr.Textbox(label="spaCy", lines=12)
        hid = gr.Textbox(visible=False)
        b1.click(task_tokenize, text_in, [s1, n1, sp1, hid])

    with gr.Tab("2. Stop Words Removal"):
        with gr.Row():
            rp = gr.Checkbox(True, label="Also remove punctuation")
            lc = gr.Checkbox(False, label="Lowercase result")
        b2 = gr.Button("Remove stop words", variant="primary")
        s2 = gr.Textbox(label="Summary")
        t2 = gr.Textbox(label="All tokens", lines=3)
        r2 = gr.Textbox(label="Removed (stop words / punctuation)", lines=3)
        f2 = gr.Textbox(label="Filtered sentence", lines=3)
        b2.click(task_stopwords, [text_in, rp, lc], [s2, t2, r2, f2])

    with gr.Tab("3. Stemming"):
        b3 = gr.Button("Apply stemming", variant="primary")
        o3 = gr.Dataframe(label="Stemming results (Porter / Lancaster / Snowball)")
        b3.click(task_stemming, text_in, o3)

    with gr.Tab("4. Lemmatization"):
        b4 = gr.Button("Apply lemmatization", variant="primary")
        o4a = gr.Dataframe(label="NLTK WordNet Lemmatizer")
        o4b = gr.Dataframe(label="spaCy Lemmatizer")
        b4.click(task_lemmatize, text_in, [o4a, o4b])

    with gr.Tab("5. POS Tags + Word Frequency"):
        b5 = gr.Button("Analyse", variant="primary")
        with gr.Row():
            o5a = gr.Dataframe(label="Part-of-speech tags")
            o5b = gr.Dataframe(label="Word frequency (without stop words)")
        b5.click(task_pos_freq, text_in, [o5a, o5b])

    with gr.Tab("6. Full Pipeline"):
        b6 = gr.Button("Run full pipeline", variant="primary")
        p = [gr.Textbox(lines=3, show_label=False) for _ in range(5)]
        final = gr.Textbox(label="Final cleaned text", lines=2)
        b6.click(task_pipeline, text_in, p + [final])

    with gr.Tab("📖 Learn"):
        gr.Markdown(LEARN_MD)

if __name__ == "__main__":
    demo.launch()