import gc
import html
import json
import re
from urllib.parse import unquote, urlparse

import gradio as gr
import numpy as np
import plotly.graph_objects as go
import requests
import spaces
import torch

from sentence_transformers import SentenceTransformer
from transformers import (
    AutoModelForCausalLM,
    AutoModelForTokenClassification,
    AutoTokenizer,
    pipeline,
)

from combo import COMBO


# ============================================================
# MODELE
# ============================================================

STELLA_MODEL = "sdadas/stella-pl-retrieval-8k"

NER_MODEL = "pczarnik/herbert-base-ner"

FASTPDN_MODEL = "clarin-pl/FastPDN"

COMBO_MODEL = (
    "clarin-pl/"
    "combo-nlp-xlm-roberta-base-polish-pdb-ud2.17"
)

SENTIMENT_MODEL = (
    "nie3e/sentiment-polish-gpt2-small"
)

LLM_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"

WIKIPEDIA_API = (
    "https://pl.wikipedia.org/w/api.php"
)

MAX_ARTICLES = 20


# ============================================================
# DOMYŚLNE URL
# ============================================================

DEFAULT_URLS = """https://pl.wikipedia.org/wiki/Barczewo_(gmina)
https://pl.wikipedia.org/wiki/Olsztyn
https://pl.wikipedia.org/wiki/Jaros%C5%82aw_Iwaszkiewicz
https://pl.wikipedia.org/wiki/Loudun
https://pl.wikipedia.org/wiki/Cystersi
https://pl.wikipedia.org/wiki/Ministerstwo_Bezpiecze%C5%84stwa_Publicznego
https://pl.wikipedia.org/wiki/Anatolij_Kaszpirowski
https://pl.wikipedia.org/wiki/Turbina_gazowa
https://pl.wikipedia.org/wiki/Fruczak_go%C5%82%C4%85bek
https://pl.wikipedia.org/wiki/Kwas_3-hydroksymas%C5%82owy"""


# ============================================================
# INFORMACJE O MODELACH
# ============================================================

MODEL_INFO = {
    "Podobieństwo semantyczne": (
        "**Model:** `sdadas/stella-pl-retrieval-8k`  \n"
        "**Zadanie:** embeddingi i podobieństwo semantyczne  \n"
        "**Metoda:** embedding → cosine similarity"
    ),

    "Mapa semantyczna 3D": (
        "**Model:** `sdadas/stella-pl-retrieval-8k`  \n"
        "**Zadanie:** reprezentacja semantyczna tekstu  \n"
        "**Metoda:** embedding → PCA → 3D"
    ),

    "NER": (
        "**Model:** `pczarnik/herbert-base-ner`  \n"
        "**Zadanie:** rozpoznawanie nazwanych encji (NER)  \n"
        "**Wynik:** typ encji, confidence, artykuł, "
        "zdanie i kontekst"
    ),

    "NER — CLARIN-PL FastPDN": (
        "**Model:** `clarin-pl/FastPDN`  \n"
        "**Zadanie:** rozpoznawanie nazwanych encji (NER)  \n"
        "**Wynik:** szczegółowy typ encji, confidence, "
        "artykuł, zdanie i kontekst"
    ),

    "COMBO — morfoskładnia i składnia": (
        "**Model:** "
        "`clarin-pl/"
        "combo-nlp-xlm-roberta-base-polish-pdb-ud2.17`  \n"
        "**Zadania:** POS, morfologia, lematyzacja "
        "i dependency parsing"
    ),

    "Sentiment — PolEmo2": (
        "**Model:** `nie3e/sentiment-polish-gpt2-small`  \n"
        "**Zadanie:** klasyfikacja sentymentu zdań  \n"
        "**Korpus:** `clarin-pl/polemo2-official`  \n"
        "**Klasy:** NEUTRAL, NEGATIVE, POSITIVE, AMBIGUOUS"
    ),

    "Podsumowanie": (
        "**Model:** `Qwen/Qwen2.5-1.5B-Instruct`  \n"
        "**Zadanie:** generowanie podsumowania"
    ),

    "JSON": (
        "**Model:** `Qwen/Qwen2.5-1.5B-Instruct`  \n"
        "**Zadanie:** ekstrakcja informacji do JSON"
    ),
}


# ============================================================
# PROMPTY
# ============================================================

PROMPTS = {
    "Podobieństwo semantyczne": (
        "Porównaj artykuły pod względem podobieństwa "
        "semantycznego na podstawie ich reprezentacji "
        "embeddingowej. Wynik przedstaw w uporządkowanej "
        "tabeli oraz wskaż najbardziej podobne pary."
    ),

    "Mapa semantyczna 3D": (
        "Przedstaw artykuły na trójwymiarowej mapie "
        "semantycznej. Położenie punktów powinno wynikać "
        "z reprezentacji embeddingowej artykułów."
    ),

    "NER": (
        "Przeprowadź uporządkowaną analizę nazwanych encji. "
        "Dla każdej rozpoznanej encji podaj: dokładny tekst "
        "encji, typ encji, confidence modelu NER, tytuł "
        "artykułu, URL źródłowy, pełne zdanie zawierające "
        "encję oraz krótki kontekst obejmujący trzy słowa "
        "przed encją i trzy słowa po niej. Nie łącz confidence "
        "NER z żadną inną miarą prawdopodobieństwa."
    ),

    "NER — CLARIN-PL FastPDN": (
        "Przeprowadź uporządkowaną analizę nazwanych encji "
        "z zachowaniem szczegółowych typów encji zwracanych "
        "przez model FastPDN. Dla każdej encji podaj: dokładny "
        "tekst encji, szczegółowy typ encji, confidence modelu "
        "NER, tytuł artykułu, URL źródłowy, pełne zdanie "
        "zawierające encję oraz krótki kontekst obejmujący "
        "trzy słowa przed encją i trzy słowa po niej."
    ),

    "COMBO — morfoskładnia i składnia": (
        "Przeprowadź uporządkowaną analizę morfosyntaktyczną "
        "i składniową. Dla każdego tokenu podaj formę, lemat, "
        "UPOS, XPOS, cechy morfologiczne, numer głowy "
        "składniowej oraz relację dependency. Podsumuj "
        "także rozkład części mowy i relacji składniowych."
    ),

    "Sentiment — PolEmo2": (
        "Przeanalizuj sentyment kolejnych zdań artykułów. "
        "Dla każdego zdania podaj: tekst zdania, klasę "
        "sentymentu, prawdopodobieństwo rozpoznanej klasy "
        "oraz wartość sentymentu w zakresie od -1 do +1. "
        "Użyj klas NEUTRAL, NEGATIVE, POSITIVE i AMBIGUOUS. "
        "POSITIVE ma wartość dodatnią, NEGATIVE ujemną, "
        "natomiast NEUTRAL i AMBIGUOUS mają wartość 0. "
        "Nie interpretuj wartości jako dodatkowego "
        "prawdopodobieństwa modelu."
    ),

    "Podsumowanie": (
        "Przygotuj neutralne, rzeczowe podsumowanie "
        "poniższego artykułu z polskiej Wikipedii. "
        "Nie dodawaj informacji, których nie ma w tekście. "
        "Podsumowanie powinno mieć 5–8 zdań."
    ),

    "JSON": (
        "Przeanalizuj poniższy artykuł z polskiej Wikipedii "
        "i zwróć wyłącznie poprawny obiekt JSON. "
        "Nie dodawaj żadnego tekstu przed JSON-em ani po nim. "
        "Nie wymyślaj informacji."
    ),
}


# ============================================================
# KOLORY SENTYMENTU
# ============================================================

SENTIMENT_STYLES = {
    "POSITIVE": {
        "background": "#eaf7ea",
        "border": "#4caf50",
        "label": "pozytywny",
    },
    "NEGATIVE": {
        "background": "#fdeaea",
        "border": "#d9534f",
        "label": "negatywny",
    },
    "NEUTRAL": {
        "background": "#fff8dc",
        "border": "#d6b656",
        "label": "neutralny",
    },
    "AMBIGUOUS": {
        "background": "#f1eafa",
        "border": "#9068b5",
        "label": "niejednoznaczny",
    },
}


SENTIMENT_SIGNS = {
    "POSITIVE": 1,
    "NEGATIVE": -1,
    "NEUTRAL": 0,
    "AMBIGUOUS": 0,
}


def calculate_sentiment_value(
    label,
    probability,
):
    sign = SENTIMENT_SIGNS.get(
        label,
        0,
    )

    return sign * probability


# ============================================================
# POMOCNICZE
# ============================================================

def cleanup_memory():
    gc.collect()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def parse_urls(text):
    urls = []

    for line in text.splitlines():
        line = line.strip()

        if line:
            urls.append(line)

    return urls[:MAX_ARTICLES]


# ============================================================
# WIKIPEDIA
# ============================================================

def extract_wikipedia_title(url):
    parsed = urlparse(url.strip())

    path = unquote(
        parsed.path
    )

    match = re.search(
        r"/wiki/(.+)$",
        path,
    )

    if not match:
        return None

    title = (
        match.group(1)
        .replace("_", " ")
        .strip()
    )

    return title or None


def fetch_article(url):
    title = extract_wikipedia_title(
        url
    )

    if not title:
        raise ValueError(
            f"Nieprawidłowy adres Wikipedii: {url}"
        )

    params = {
        "action": "query",
        "format": "json",
        "prop": "extracts|info",
        "explaintext": 1,
        "inprop": "url",
        "titles": title,
    }

    response = requests.get(
        WIKIPEDIA_API,
        params=params,
        timeout=30,
        headers={
            "User-Agent": (
                "Wikipedia-NLP-Experiment/1.0"
            )
        },
    )

    response.raise_for_status()

    data = response.json()

    pages = (
        data
        .get("query", {})
        .get("pages", {})
    )

    if not pages:
        raise ValueError(
            f"Nie znaleziono artykułu: {title}"
        )

    page = next(
        iter(pages.values())
    )

    if "missing" in page:
        raise ValueError(
            f"Nie znaleziono artykułu: {title}"
        )

    return {
        "title": page.get(
            "title",
            title,
        ),
        "url": page.get(
            "fullurl",
            url,
        ),
        "text": page.get(
            "extract",
            "",
        ),
    }


def get_stats(article):
    text = article["text"]

    paragraphs = [
        p.strip()
        for p in re.split(
            r"\n\s*\n",
            text,
        )
        if p.strip()
    ]

    sentences = [
        s.strip()
        for s in re.split(
            r"(?<=[.!?])\s+",
            text,
        )
        if s.strip()
    ]

    words = re.findall(
        r"\b[\wÀ-ž'-]+\b",
        text,
        flags=re.UNICODE,
    )

    return {
        "znaki": len(text),
        "slowa": len(words),
        "akapity": len(paragraphs),
        "zdania": len(sentences),
    }


def load_articles(url_text):
    urls = parse_urls(url_text)

    if not urls:
        return (
            [],
            "Nie podano adresów URL.",
        )

    articles = []
    errors = []

    for url in urls:
        try:
            article = fetch_article(
                url
            )

            article["stats"] = get_stats(
                article
            )

            articles.append(article)

        except Exception as error:
            errors.append(
                f"{url}: {error}"
            )

    if not articles:
        return (
            [],
            "Nie udało się pobrać żadnego "
            "artykułu.\n\n"
            + "\n".join(errors),
        )

    message = (
        f"Pobrano {len(articles)} artykułów."
    )

    if errors:
        message += (
            "\n\nNie udało się pobrać:\n"
            + "\n".join(errors)
        )

    return articles, message


def articles_table(articles):
    rows = []

    for article in articles:
        stats = article["stats"]

        rows.append(
            [
                article["title"],
                stats["slowa"],
                stats["znaki"],
                stats["akapity"],
                stats["zdania"],
                article["url"],
            ]
        )

    return rows


# ============================================================
# STELLA
# ============================================================

@spaces.GPU(duration=180)
def run_stella(articles):
    if len(articles) < 2:
        return (
            "Do analizy podobieństwa potrzebne "
            "są co najmniej 2 artykuły.",
            None,
        )

    device = (
        "cuda"
        if torch.cuda.is_available()
        else "cpu"
    )

    model = SentenceTransformer(
        STELLA_MODEL,
        trust_remote_code=True,
        device=device,
    )

    texts = []

    for article in articles:
        text = article["text"]

        if len(text) > 30000:
            text = text[:30000]

        texts.append(text)

    embeddings = model.encode(
        texts,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )

    similarity = np.matmul(
        embeddings,
        embeddings.T,
    )

    lines = [
        "## Podobieństwo semantyczne",
        "",
        "Wartości oznaczają podobieństwo "
        "cosinusowe embeddingów.",
        "",
        "| Artykuł | "
        + " | ".join(
            article["title"]
            for article in articles
        )
        + " |",
        "|---|"
        + "|".join(
            ["---"] * len(articles)
        )
        + "|",
    ]

    for i, article in enumerate(
        articles
    ):
        values = [
            f"{similarity[i, j]:.3f}"
            for j in range(
                len(articles)
            )
        ]

        lines.append(
            "| "
            + article["title"]
            + " | "
            + " | ".join(values)
            + " |"
        )

    pairs = []

    for i in range(
        len(articles)
    ):
        for j in range(
            i + 1,
            len(articles),
        ):
            pairs.append(
                (
                    similarity[i, j],
                    articles[i]["title"],
                    articles[j]["title"],
                )
            )

    pairs.sort(
        reverse=True
    )

    lines.extend(
        [
            "",
            "### Najbardziej podobne pary",
            "",
        ]
    )

    for score, first, second in pairs[:10]:
        lines.append(
            f"- **{first}** ↔ "
            f"**{second}**: {score:.3f}"
        )

    del model
    del embeddings

    cleanup_memory()

    return (
        "\n".join(lines),
        None,
    )


# ============================================================
# MAPA 3D
# ============================================================

def create_3d_coordinates(
    embeddings
):
    data = np.asarray(
        embeddings,
        dtype=np.float64,
    )

    centered = (
        data - data.mean(axis=0)
    )

    if len(data) < 2:
        return np.zeros(
            (len(data), 3)
        )

    covariance = np.cov(
        centered,
        rowvar=False,
    )

    eigenvalues, eigenvectors = np.linalg.eigh(
        covariance
    )

    order = np.argsort(
        eigenvalues
    )[::-1]

    eigenvectors = eigenvectors[
        :,
        order,
    ]

    dimensions = min(
        3,
        eigenvectors.shape[1],
    )

    coordinates = (
        centered
        @ eigenvectors[:, :dimensions]
    )

    if dimensions < 3:
        coordinates = np.pad(
            coordinates,
            (
                (0, 0),
                (0, 3 - dimensions),
            ),
        )

    return coordinates


@spaces.GPU(duration=180)
def run_semantic_map(articles):
    if len(articles) < 2:
        return (
            "Do stworzenia mapy potrzebne "
            "są co najmniej 2 artykuły.",
            None,
        )

    device = (
        "cuda"
        if torch.cuda.is_available()
        else "cpu"
    )

    model = SentenceTransformer(
        STELLA_MODEL,
        trust_remote_code=True,
        device=device,
    )

    texts = []

    for article in articles:
        text = article["text"]

        if len(text) > 30000:
            text = text[:30000]

        texts.append(text)

    embeddings = model.encode(
        texts,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )

    coordinates = create_3d_coordinates(
        embeddings
    )

    figure = go.Figure()

    figure.add_trace(
        go.Scatter3d(
            x=coordinates[:, 0],
            y=coordinates[:, 1],
            z=coordinates[:, 2],
            mode="markers+text",
            text=[
                article["title"]
                for article in articles
            ],
            textposition="top center",
            hovertext=[
                article["url"]
                for article in articles
            ],
            hoverinfo="text",
            marker={
                "size": 8,
                "opacity": 0.85,
            },
        )
    )

    figure.update_layout(
        title="Mapa semantyczna 3D",
        height=700,
        scene={
            "xaxis_title": "PCA 1",
            "yaxis_title": "PCA 2",
            "zaxis_title": "PCA 3",
        },
        margin={
            "l": 0,
            "r": 0,
            "t": 50,
            "b": 0,
        },
    )

    del model
    del embeddings

    cleanup_memory()

    return (
        "## Mapa semantyczna 3D",
        figure,
    )


# ============================================================
# NER — POMOCNICZE
# ============================================================

def split_for_ner(
    text,
    max_chars=1800,
):
    if len(text) <= max_chars:
        return [text]

    sentences = re.split(
        r"(?<=[.!?])\s+",
        text,
    )

    chunks = []
    current = ""

    for sentence in sentences:
        sentence = sentence.strip()

        if not sentence:
            continue

        candidate = (
            sentence
            if not current
            else current + " " + sentence
        )

        if len(candidate) > max_chars:
            if current:
                chunks.append(
                    current
                )

            current = sentence

        else:
            current = candidate

    if current:
        chunks.append(
            current
        )

    return chunks


def get_sentence_spans(text):
    spans = []

    pattern = re.compile(
        r".+?(?:[.!?](?=\s|$)|$)",
        flags=re.DOTALL,
    )

    for match in pattern.finditer(text):
        start = match.start()
        end = match.end()

        while (
            start < end
            and text[start].isspace()
        ):
            start += 1

        while (
            end > start
            and text[end - 1].isspace()
        ):
            end -= 1

        if end > start:
            spans.append(
                {
                    "start": start,
                    "end": end,
                    "text": text[start:end],
                }
            )

    return spans


def find_sentence_for_entity(
    sentence_spans,
    start,
    end,
):
    if start is None or end is None:
        return None

    for sentence in sentence_spans:
        if (
            sentence["start"] <= start
            and sentence["end"] >= end
        ):
            return sentence

    return None


def get_word_context(
    text,
    start,
    end,
    before=3,
    after=3,
):
    if start is None or end is None:
        return text

    word_spans = list(
        re.finditer(
            r"\b[\wÀ-ž'-]+\b",
            text,
            flags=re.UNICODE,
        )
    )

    entity_index = None

    for index, match in enumerate(
        word_spans
    ):
        if (
            match.start() <= start
            and match.end() >= end
        ):
            entity_index = index
            break

    if entity_index is None:
        return text

    first = max(
        0,
        entity_index - before,
    )

    last = min(
        len(word_spans),
        entity_index + after + 1,
    )

    return " ".join(
        match.group(0)
        for match in word_spans[
            first:last
        ]
    )


# ============================================================
# NER
# ============================================================

@spaces.GPU(duration=300)
def run_ner_model(
    articles,
    model_name,
):
    tokenizer = AutoTokenizer.from_pretrained(
        model_name
    )

    model = AutoModelForTokenClassification.from_pretrained(
        model_name
    )

    device = (
        0
        if torch.cuda.is_available()
        else -1
    )

    ner = pipeline(
        "token-classification",
        model=model,
        tokenizer=tokenizer,
        aggregation_strategy="simple",
        device=device,
    )

    rows = []

    for article in articles:
        chunks = split_for_ner(
            article["text"]
        )

        for chunk in chunks:
            sentence_spans = (
                get_sentence_spans(
                    chunk
                )
            )

            results = ner(chunk)

            for result in results:
                entity_text = (
                    result
                    .get("word", "")
                    .strip()
                )

                if not entity_text:
                    continue

                start = result.get(
                    "start"
                )

                end = result.get(
                    "end"
                )

                sentence = (
                    find_sentence_for_entity(
                        sentence_spans,
                        start,
                        end,
                    )
                )

                entity_type = result.get(
                    "entity_group",
                    result.get(
                        "entity",
                        "UNKNOWN",
                    ),
                )

                confidence = float(
                    result.get(
                        "score",
                        0,
                    )
                )

                # Tylko krótki kontekst:
                # 3 słowa przed + encja + 3 słowa po.
                context = get_word_context(
                    chunk,
                    start,
                    end,
                    before=3,
                    after=3,
                )

                rows.append(
                    [
                        entity_type,
                        entity_text,
                        article["title"],
                        article["url"],
                        (
                            sentence["text"]
                            if sentence
                            else ""
                        ),
                        context,
                        confidence,
                    ]
                )

    del ner
    del model
    del tokenizer

    cleanup_memory()

    # NIE TWORZYMY tutaj tabeli Markdown.
    # Wyniki są wyświetlane wyłącznie przez
    # osobną tabelę gr.Dataframe.
    return (
        f"## Rozpoznane encje\n\n"
        f"**Model:** `{model_name}`\n\n"
        f"Znaleziono **{len(rows)}** "
        f"wystąpień encji.",
        rows,
    )


# ============================================================
# SENTIMENT
# ============================================================

def split_into_sentences(text):
    sentences = re.split(
        r"(?<=[.!?])\s+",
        text,
    )

    return [
        sentence.strip()
        for sentence in sentences
        if sentence.strip()
    ]


def sentiment_label(
    label
):
    return label.upper().strip()


def build_sentiment_html(
    article_results
):
    parts = [
        """
        <style>
            .sentiment-container {
                font-family:
                    system-ui,
                    -apple-system,
                    BlinkMacSystemFont,
                    "Segoe UI",
                    sans-serif;
                line-height: 1.55;
            }

            .sentiment-legend {
                display: flex;
                flex-wrap: wrap;
                gap: 8px;
                margin-bottom: 20px;
            }

            .sentiment-legend span {
                padding: 5px 10px;
                border-radius: 999px;
                font-size: .82rem;
                border: 1px solid rgba(0,0,0,.12);
            }

            .sentiment-article {
                margin-bottom: 30px;
            }

            .sentiment-title {
                font-size: 1.15rem;
                font-weight: 700;
                margin-bottom: 4px;
            }

            .sentiment-url {
                font-size: .78rem;
                opacity: .65;
                margin-bottom: 12px;
                word-break: break-all;
            }

            .sentiment-sentence {
                border-left: 5px solid;
                border-radius: 6px;
                padding: 12px 15px;
                margin: 9px 0;
            }

            .sentiment-meta {
                display: flex;
                flex-wrap: wrap;
                gap: 7px;
                margin-bottom: 8px;
            }

            .sentiment-badge {
                padding: 3px 8px;
                border-radius: 999px;
                background: rgba(255,255,255,.7);
                border: 1px solid rgba(0,0,0,.12);
                font-size: .78rem;
                font-weight: 600;
            }

            .sentiment-text {
                font-size: .98rem;
            }
        </style>

        <div class="sentiment-container">

            <div class="sentiment-legend">
                <span style="
                    background:#eaf7ea;
                    border-color:#4caf50;
                ">
                    ● pozytywny
                </span>

                <span style="
                    background:#fdeaea;
                    border-color:#d9534f;
                ">
                    ● negatywny
                </span>

                <span style="
                    background:#fff8dc;
                    border-color:#d6b656;
                ">
                    ● neutralny
                </span>

                <span style="
                    background:#f1eafa;
                    border-color:#9068b5;
                ">
                    ● niejednoznaczny
                </span>
            </div>
        """
    ]

    for article in article_results:
        parts.extend(
            [
                "<div class='sentiment-article'>",

                "<div class='sentiment-title'>"
                + html.escape(
                    article["title"]
                )
                + "</div>",

                "<div class='sentiment-url'>"
                + html.escape(
                    article["url"]
                )
                + "</div>",
            ]
        )

        for item in article["sentences"]:
            label = item["label"]
            probability = item["probability"]
            value = item["value"]

            style = SENTIMENT_STYLES.get(
                label,
                SENTIMENT_STYLES[
                    "NEUTRAL"
                ],
            )

            parts.extend(
                [
                    "<div "
                    "class='sentiment-sentence' "
                    f"style='background:"
                    f"{style['background']};"
                    f"border-color:"
                    f"{style['border']};'>",

                    "<div "
                    "class='sentiment-meta'>",

                    "<span "
                    "class='sentiment-badge'>"
                    + style["label"]
                    + "</span>",

                    "<span "
                    "class='sentiment-badge'>"
                    "prawdopodobieństwo: "
                    f"{probability:.3f}"
                    "</span>",

                    "<span "
                    "class='sentiment-badge'>"
                    "wartość: "
                    f"{value:+.3f}"
                    "</span>",

                    "</div>",

                    "<div "
                    "class='sentiment-text'>"
                    + html.escape(
                        item["sentence"]
                    )
                    + "</div>",

                    "</div>",
                ]
            )

        parts.append(
            "</div>"
        )

    parts.append(
        "</div>"
    )

    return "".join(parts)


@spaces.GPU(duration=300)
def run_sentiment(articles):
    device = (
        0
        if torch.cuda.is_available()
        else -1
    )

    sentiment = pipeline(
        "text-classification",
        model=SENTIMENT_MODEL,
        tokenizer=SENTIMENT_MODEL,
        device=device,
    )

    table_rows = []
    article_results = []

    total_counts = {
        "POSITIVE": 0,
        "NEGATIVE": 0,
        "NEUTRAL": 0,
        "AMBIGUOUS": 0,
    }

    for article in articles:
        sentences = split_into_sentences(
            article["text"]
        )

        sentences = [
            sentence
            for sentence in sentences
            if len(sentence) >= 10
        ]

        if not sentences:
            continue

        results = sentiment(
            sentences,
            batch_size=16,
            truncation=True,
            max_length=512,
        )

        current_article = {
            "title": article["title"],
            "url": article["url"],
            "sentences": [],
        }

        for sentence, result in zip(
            sentences,
            results,
        ):
            label = sentiment_label(
                result["label"]
            )

            probability = float(
                result["score"]
            )

            value = calculate_sentiment_value(
                label,
                probability,
            )

            if label not in total_counts:
                total_counts[label] = 0

            total_counts[label] += 1

            current_article[
                "sentences"
            ].append(
                {
                    "sentence": sentence,
                    "label": label,
                    "probability": probability,
                    "value": value,
                }
            )

            table_rows.append(
                [
                    article["title"],
                    sentence,
                    label,
                    probability,
                    value,
                ]
            )

        article_results.append(
            current_article
        )

    del sentiment

    cleanup_memory()

    lines = [
        "## Analiza sentymentu",
        "",
        f"**Model:** `{SENTIMENT_MODEL}`",
        "",
        "### Klasy",
        "",
        "| Klasa | Liczba zdań |",
        "|---|---:|",
    ]

    for label in [
        "POSITIVE",
        "NEGATIVE",
        "NEUTRAL",
        "AMBIGUOUS",
    ]:
        lines.append(
            f"| {label} | "
            f"{total_counts.get(label, 0)} |"
        )

    lines.extend(
        [
            "",
            "### Interpretacja wartości",
            "",
            "- `POSITIVE` → wartość dodatnia",
            "- `NEGATIVE` → wartość ujemna",
            "- `NEUTRAL` → `0`",
            "- `AMBIGUOUS` → `0`",
            "",
            "Wartość jest iloczynem znaku klasy "
            "i prawdopodobieństwa klasy.",
            "Nie jest to osobna predykcja modelu.",
        ]
    )

    return (
        "\n".join(lines),
        build_sentiment_html(
            article_results
        ),
        table_rows,
    )


# ============================================================
# COMBO
# ============================================================

@spaces.GPU(duration=300)
def run_combo(articles):
    nlp = COMBO.from_pretrained(
        COMBO_MODEL
    )

    rows = []
    pos_counts = {}
    relation_counts = {}

    for article in articles:
        text = article["text"]

        if len(text) > 20000:
            text = text[:20000]

        result = nlp(text)

        for sentence in result:
            for token in sentence:
                pos = token.upos
                relation = token.deprel

                pos_counts[pos] = (
                    pos_counts.get(
                        pos,
                        0,
                    )
                    + 1
                )

                relation_counts[relation] = (
                    relation_counts.get(
                        relation,
                        0,
                    )
                    + 1
                )

                rows.append(
                    [
                        article["title"],
                        token.form,
                        token.lemma,
                        token.upos,
                        token.xpos,
                        str(token.feats),
                        token.head,
                        token.deprel,
                    ]
                )

    del nlp

    cleanup_memory()

    lines = [
        "## Analiza morfosyntaktyczna "
        "i składniowa",
        "",
        f"**Model:** `{COMBO_MODEL}`",
        "",
        "### Części mowy",
        "",
        "| UPOS | Liczba |",
        "|---|---:|",
    ]

    for pos, count in sorted(
        pos_counts.items(),
        key=lambda item: -item[1],
    ):
        lines.append(
            f"| {pos} | {count} |"
        )

    lines.extend(
        [
            "",
            "### Relacje składniowe",
            "",
            "| Relacja | Liczba |",
            "|---|---:|",
        ]
    )

    for relation, count in sorted(
        relation_counts.items(),
        key=lambda item: -item[1],
    ):
        lines.append(
            f"| {relation} | {count} |"
        )

    lines.extend(
        [
            "",
            "### Tokeny",
            "",
            "| Artykuł | Forma | Lemma | "
            "UPOS | XPOS | Cechy | "
            "Head | Relacja |",
            "|---|---|---|---|---|---|---:|---|",
        ]
    )

    for row in rows[:5000]:
        lines.append(
            "| "
            + " | ".join(
                str(value)
                .replace("|", "\\|")
                .replace("\n", " ")
                for value in row
            )
            + " |"
        )

    return "\n".join(lines)


# ============================================================
# LLM
# ============================================================

def build_llm_prompt(
    prompt,
    article,
):
    return (
        prompt
        + "\n\n"
        + "ARTYKUŁ:\n\n"
        + article["text"]
    )


@spaces.GPU(duration=300)
def run_llm(
    articles,
    prompt,
    max_tokens,
    mode,
):
    tokenizer = AutoTokenizer.from_pretrained(
        LLM_MODEL
    )

    model = AutoModelForCausalLM.from_pretrained(
        LLM_MODEL,
        torch_dtype=(
            torch.float16
            if torch.cuda.is_available()
            else torch.float32
        ),
        device_map=(
            "auto"
            if torch.cuda.is_available()
            else None
        ),
    )

    outputs = []

    for article in articles:
        text = article["text"]

        if len(text) > 30000:
            text = text[:30000]

        messages = [
            {
                "role": "system",
                "content": (
                    "Jesteś pomocnym analitykiem "
                    "tekstów w języku polskim."
                ),
            },
            {
                "role": "user",
                "content": (
                    prompt
                    + "\n\n"
                    + "ARTYKUŁ:\n\n"
                    + text
                ),
            },
        ]

        formatted = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

        inputs = tokenizer(
            formatted,
            return_tensors="pt",
            truncation=True,
            max_length=8192,
        )

        if torch.cuda.is_available():
            inputs = {
                key: value.to(
                    model.device
                )
                for key, value in inputs.items()
            }

        with torch.no_grad():
            generated = model.generate(
                **inputs,
                max_new_tokens=int(
                    max_tokens
                ),
                do_sample=False,
                temperature=None,
                pad_token_id=(
                    tokenizer.eos_token_id
                ),
            )

        generated = generated[
            0,
            inputs["input_ids"].shape[1]:,
        ]

        output = tokenizer.decode(
            generated,
            skip_special_tokens=True,
        ).strip()

        outputs.append(
            "## "
            + article["title"]
            + "\n\n"
            + output
        )

    del model
    del tokenizer

    cleanup_memory()

    return "\n\n---\n\n".join(
        outputs
    )


# ============================================================
# JSON
# ============================================================

def try_parse_json(text):
    text = text.strip()

    try:
        return json.loads(text)

    except json.JSONDecodeError:
        match = re.search(
            r"\{.*\}",
            text,
            flags=re.DOTALL,
        )

        if match:
            try:
                return json.loads(
                    match.group(0)
                )
            except json.JSONDecodeError:
                pass

    return None


# ============================================================
# ANALIZA GŁÓWNA
# ============================================================

def analyze(
    analysis,
    articles,
    max_tokens,
    prompt,
):
    if not articles:
        return (
            "Najpierw pobierz artykuły.",
            None,
            [],
            "",
        )

    if not prompt.strip():
        prompt = PROMPTS.get(
            analysis,
            "",
        )

    if analysis == "Podobieństwo semantyczne":
        output, plot = run_stella(
            articles
        )

        return (
            output,
            plot,
            [],
            "",
        )

    if analysis == "Mapa semantyczna 3D":
        output, plot = run_semantic_map(
            articles
        )

        return (
            output,
            plot,
            [],
            "",
        )

    if analysis == "NER":
        output, rows = run_ner_model(
            articles,
            NER_MODEL,
        )

        return (
            output,
            None,
            rows,
            "",
        )

    if analysis == "NER — CLARIN-PL FastPDN":
        output, rows = run_ner_model(
            articles,
            FASTPDN_MODEL,
        )

        return (
            output,
            None,
            rows,
            "",
        )

    if analysis == (
        "COMBO — morfoskładnia i składnia"
    ):
        output = run_combo(
            articles
        )

        return (
            output,
            None,
            [],
            "",
        )

    if analysis == "Sentiment — PolEmo2":
        output, sentiment_html, rows = (
            run_sentiment(
                articles
            )
        )

        return (
            output,
            None,
            [],
            sentiment_html,
        )

    if analysis in (
        "Podsumowanie",
        "JSON",
    ):
        output = run_llm(
            articles,
            prompt,
            max_tokens,
            analysis,
        )

        return (
            output,
            None,
            [],
            "",
        )

    return (
        "Nieznany typ analizy.",
        None,
        [],
        "",
    )


# ============================================================
# UI
# ============================================================

ANALYSES = [
    "Podobieństwo semantyczne",
    "Mapa semantyczna 3D",
    "NER",
    "NER — CLARIN-PL FastPDN",
    "COMBO — morfoskładnia i składnia",
    "Sentiment — PolEmo2",
    "Podsumowanie",
    "JSON",
]


def analysis_changed(
    analysis
):
    return (
        MODEL_INFO.get(
            analysis,
            "",
        ),
        PROMPTS.get(
            analysis,
            "",
        ),
    )


with gr.Blocks(
    title="Wikipedia NLP",
) as demo:

    gr.Markdown(
        """
# Wikipedia NLP

Analiza artykułów polskiej Wikipedii z wykorzystaniem
modeli semantycznych, NER, morfoskładni, sentimentu
oraz modeli generatywnych.
"""
    )

    # --------------------------------------------------------
    # URL
    # --------------------------------------------------------

    with gr.Row():

        with gr.Column(
            scale=3
        ):
            urls = gr.Textbox(
                label="Adresy artykułów Wikipedii",
                value=DEFAULT_URLS,
                lines=12,
                placeholder=(
                    "Jeden adres URL w każdym wierszu"
                ),
            )

        with gr.Column(
            scale=1
        ):
            load_button = gr.Button(
                "Pobierz artykuły",
                variant="primary",
            )

            load_status = gr.Markdown(
                "Nie pobrano jeszcze artykułów."
            )

    # --------------------------------------------------------
    # TABELA ARTYKUŁÓW
    # --------------------------------------------------------

    article_table = gr.Dataframe(
        headers=[
            "Artykuł",
            "Słowa",
            "Znaki",
            "Akapity",
            "Zdania",
            "URL",
        ],
        datatype=[
            "str",
            "number",
            "number",
            "number",
            "number",
            "str",
        ],
        label="Pobrane artykuły",
        interactive=False,
        wrap=True,
    )

    articles_state = gr.State([])

    # --------------------------------------------------------
    # ANALIZA
    # --------------------------------------------------------

    gr.Markdown(
        "## Analiza"
    )

    with gr.Row():

        with gr.Column(
            scale=2
        ):
            analysis_dropdown = gr.Dropdown(
                choices=ANALYSES,
                value="NER",
                label="Rodzaj analizy",
            )

        with gr.Column(
            scale=3
        ):
            model_info = gr.Markdown(
                MODEL_INFO["NER"]
            )

    prompt = gr.Textbox(
        label="Prompt",
        value=PROMPTS["NER"],
        lines=8,
        interactive=True,
    )

    with gr.Row():

        max_tokens = gr.Slider(
            minimum=128,
            maximum=4096,
            value=768,
            step=128,
            label="Maksymalna liczba tokenów "
                  "dla modeli generatywnych",
        )

        run_button = gr.Button(
            "Uruchom analizę",
            variant="primary",
        )

    # --------------------------------------------------------
    # WYNIK
    # --------------------------------------------------------

    output = gr.Markdown(
        label="Wynik"
    )

    # --------------------------------------------------------
    # SENTIMENT — KOLOROWANE ZDANIA
    # --------------------------------------------------------

    sentiment_html = gr.HTML(
        label="Zdania z oznaczeniem sentymentu",
        visible=True,
    )

    # --------------------------------------------------------
    # NER — JEDYNA TABELA
    # --------------------------------------------------------

    ner_table = gr.Dataframe(
        headers=[
            "Typ",
            "Encja",
            "Artykuł",
            "URL",
            "Zdanie",
            "Kontekst",
            "Confidence",
        ],
        datatype=[
            "str",
            "str",
            "str",
            "str",
            "str",
            "str",
            "number",
        ],
        label="Wyniki NER",
        interactive=False,
        wrap=True,
        visible=True,
    )

    # --------------------------------------------------------
    # WYKRES
    # --------------------------------------------------------

    plot = gr.Plot(
        label="Wizualizacja"
    )

    # --------------------------------------------------------
    # ZMIANA ANALIZY
    # --------------------------------------------------------

    analysis_dropdown.change(
        fn=analysis_changed,
        inputs=analysis_dropdown,
        outputs=[
            model_info,
            prompt,
        ],
    )

    # --------------------------------------------------------
    # POBIERANIE
    # --------------------------------------------------------

    load_button.click(
        fn=load_articles,
        inputs=urls,
        outputs=[
            articles_state,
            load_status,
        ],
    ).then(
        fn=articles_table,
        inputs=articles_state,
        outputs=article_table,
    )

    # --------------------------------------------------------
    # URUCHOMIENIE ANALIZY
    # --------------------------------------------------------

    run_button.click(
        fn=analyze,
        inputs=[
            analysis_dropdown,
            articles_state,
            max_tokens,
            prompt,
        ],
        outputs=[
            output,
            plot,
            ner_table,
            sentiment_html,
        ],
    )


# ============================================================
# START
# ============================================================

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