from functools import lru_cache
from itertools import combinations
import os
import re

import gensim.downloader as api
import gradio as gr
import pandas as pd
from sentence_transformers import SentenceTransformer, util


MODEL_NAME = "Alibaba-NLP/gte-modernbert-base"
GLOVE_NAME = "glove-wiki-gigaword-50"
MAX_WORDS = 20


@lru_cache(maxsize=1)
def load_glove_model():
    return api.load(GLOVE_NAME)


@lru_cache(maxsize=1)
def load_contextual_model():
    return SentenceTransformer(MODEL_NAME)


def parse_words(words_text):
    words = [
        item.strip()
        for item in re.split(r"[\n,]+", words_text)
        if item.strip()
    ]
    return list(dict.fromkeys(words))


def get_token_span(tokens, word):
    cleaned = [
        token.replace("##", "").replace("\u0120", "").replace("\u2581", "")
        for token in tokens
    ]
    word_clean = word.lower()

    for start in range(len(cleaned)):
        current = ""
        for end in range(start, len(cleaned)):
            current += cleaned[end].lower()

            if current == word_clean:
                return list(range(start, end + 1))

            if len(current) > len(word_clean):
                break

    return None


def contextual_similarity(sent_a, sent_b, word_a, word_b, model):
    token_embeddings = model.encode(
        [sent_a, sent_b],
        output_value="token_embeddings",
        show_progress_bar=False,
    )
    tokens_a = model.tokenizer.tokenize(sent_a)
    tokens_b = model.tokenizer.tokenize(sent_b)

    span_a = get_token_span(tokens_a, word_a)
    span_b = get_token_span(tokens_b, word_b)

    if span_a is None or span_b is None:
        return None

    vec_a = token_embeddings[0][span_a].mean(dim=0)
    vec_b = token_embeddings[1][span_b].mean(dim=0)
    return util.cos_sim(vec_a, vec_b).item()


def static_similarity(word_a, word_b, glove_model):
    try:
        return glove_model.similarity(word_a, word_b)
    except KeyError:
        return None


def compare_word_list(pre_context, words_text, post_context):
    words = parse_words(words_text)
    if len(words) < 2:
        raise gr.Error("Enter at least two words, separated by commas or new lines.")
    if len(words) > MAX_WORDS:
        raise gr.Error(f"Enter {MAX_WORDS} words or fewer.")

    glove_model = load_glove_model()
    contextual_model = load_contextual_model()

    rows = []
    for word_a, word_b in combinations(words, 2):
        sent_a = f"{pre_context}{word_a}{post_context}"
        sent_b = f"{pre_context}{word_b}{post_context}"

        static_sim = static_similarity(word_a, word_b, glove_model)
        contextual_sim = contextual_similarity(
            sent_a, sent_b, word_a, word_b, contextual_model
        )

        rows.append(
            {
                "Word A": word_a,
                "Word B": word_b,
                "Sentence A": sent_a,
                "Sentence B": sent_b,
                "Static Sim": None if static_sim is None else round(static_sim, 4),
                f"Contextual Sim ({MODEL_NAME})": (
                    None if contextual_sim is None else round(contextual_sim, 4)
                ),
            }
        )

    return pd.DataFrame(rows)


with gr.Blocks(title="Static vs. Contextual Embeddings") as demo:
    gr.Markdown(
        "# Static vs. Contextual Embeddings\n"
        "Compare every pair in a word list using static GloVe similarity and "
        "contextual token similarity from ModernBERT."
    )

    with gr.Row():
        pre_context = gr.Textbox(
            label="Pre-context",
            value="She took a firm ",
            lines=2,
        )
        post_context = gr.Textbox(
            label="Post-context",
            value=" on the issue.",
            lines=2,
        )

    words = gr.Textbox(
        label="Words to compare",
        value="stance, stand, posture, position",
        lines=4,
        info="Separate words with commas or new lines.",
    )

    compare_button = gr.Button("Compare", variant="primary")
    output = gr.Dataframe(
        label="Pairwise similarities",
        wrap=True,
        interactive=False,
    )

    compare_button.click(
        fn=compare_word_list,
        inputs=[pre_context, words, post_context],
        outputs=output,
    )

    gr.Examples(
        examples=[
            ["I really enjoy eating ", "meat, fruit, steak, apples", " for lunch."],
            ["She took a firm ", "stance, stand, posture, position", " on the issue."],
            ["The camera ", "stand, stance", " collapsed during filming."],
        ],
        inputs=[pre_context, words, post_context],
    )


if __name__ == "__main__":
    port = int(os.environ.get("PORT", os.environ.get("GRADIO_SERVER_PORT", "7860")))
    demo.launch(server_name="0.0.0.0", server_port=port, ssr_mode=False)