from transformers import AutoModelForTokenClassification, AutoTokenizer
import numpy as np
import scipy.special
import torch
import gradio as gr
from torch.utils.data import Dataset
import re

# --- Маска ---
def make_last_subtoken_mask(mask, has_cls=True, has_eos=True):
    if has_cls:
        mask = mask[1:]
    if has_eos:
        mask = mask[:-1]
    is_last_word = list((first != second) for first, second in zip(mask[:-1], mask[1:])) + [True]
    if has_cls:
        is_last_word = [False] + is_last_word
    if has_eos:
        is_last_word.append(False)
    return is_last_word

# --- Класс UDDataset ---
class UDDataset(Dataset):

    def __init__(self, data, tokenizer, min_count=1, tags=None): 
        self.data = data
        self.tokenizer = tokenizer
        self.raw_labels = [item["labels"] for item in data if "labels" in item]
        if tags is None:
            tag_counts = Counter([tag for elem in data for tag in elem["labels"]])
            self.tags_ = ["<PAD>", "<UNK>"] + [x for x, count in tag_counts.items() if count >= min_count]
        else:
            self.tags_ = tags
        self.tag_indexes_ = {tag: i for i, tag in enumerate(self.tags_)}
        self.unk_index = 1 #0
        self.ignore_index = -100

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        item = self.data[index]
        tokenization = self.tokenizer(item["words"], is_split_into_words=True)
        last_subtoken_mask = make_last_subtoken_mask(tokenization.word_ids())
        answer = {"input_ids": tokenization["input_ids"], 
                  "mask": last_subtoken_mask, 
                  "attention_mask": tokenization["attention_mask"]}
        if "labels" in item:
            labels = [self.tag_indexes_.get(tag, self.unk_index) for tag in item["labels"]]
            zero_labels = np.array([self.ignore_index] * len(tokenization["input_ids"]), dtype=int)
            zero_labels[last_subtoken_mask] = labels
            answer["labels"] = zero_labels
        return answer

# --- Загрузка модели и токенизатора ---
model_name = "ossetic-encoders/ossbert-morph-v2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
id2label = model.config.id2label
classes = [id2label[i] for i in range(len(id2label))]

model_name_l = "ossetic-encoders/ossbert-lemm-v2"
tokenizer_l = AutoTokenizer.from_pretrained(model_name_l)
model_l = AutoModelForTokenClassification.from_pretrained(model_name_l)
id2label_l = model_l.config.id2label
classes_l = [id2label_l[i] for i in range(len(id2label_l))]

# --- Получение предсказаний ---
def predict_top_k(model, dataset, classes, top_k):
    model.eval()
    answer = []
    with torch.no_grad():
        for elem in dataset:
            input_ids = torch.tensor(elem["input_ids"]).unsqueeze(0)
            attention_mask = torch.tensor(elem["attention_mask"]).unsqueeze(0)

            inputs = {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
            }

            outputs = model(**inputs)
            logits = outputs.logits.squeeze().numpy()
            mask = elem["mask"]
            probs = scipy.special.softmax(logits, axis=-1)[:len(mask)]

            top_k_indices = np.argsort(probs, axis=-1)[:, -top_k:][:, ::-1] 
            top_k_probs = np.take_along_axis(probs, top_k_indices, axis=-1)

            top_k_labels = []
            for i in range(len(mask)):
                if mask[i]:
                    labels = [classes[idx] for idx in top_k_indices[i]]
                    probs = [f"{p:.2f}" for p in top_k_probs[i]]
                    top_k_labels.append([(label, prob) for label, prob in zip(labels, probs)])

            answer.append({"top_k_labels": top_k_labels})
    return answer

def restore_lemma(word_form, label):
    try:
        lemma_rule, form_rule = label.split('#')
        form_parts = form_rule.split('+')
        form_constants = [part for part in form_parts if not part.isdigit()]

        extracted_vars = {}

        regex_pattern = ""
        var_order = []
        for part in form_parts:
                if part.isdigit():
                    regex_pattern += r"(.+)"
                    var_order.append(int(part))
                else:
                    regex_pattern += re.escape(part)

        match = re.match(f"^{regex_pattern}$", word_form)
        if match:
                extracted_vars = {var_num: val for var_num, val in zip(var_order, match.groups())}
        else:
            suppl = [p for p in lemma_rule.split('+') if not p.isdigit()]
            nums = [p for p in lemma_rule.split('+') if p.isdigit()]
            if len(suppl) == 1 and len(nums) == 1:
                return suppl[0]
            else:
                return word_form
        lemma_parts = lemma_rule.split('+')
        final_lemma_pieces = []

        for part in lemma_parts:
            if part.isdigit():
                var_num = int(part)
                final_lemma_pieces.append(extracted_vars.get(var_num, ""))
            else:
                final_lemma_pieces.append(part)
        return "".join(final_lemma_pieces)

    except Exception:
        return word_form

#--- Функция для Gradio ---
def analyze_text(text, top_k_lemmas, top_k_tags, show_paradigm, show_subtokens):
    text = text.replace('ӕ', 'æ')
    text= text.replace('Ӕ', 'Æ')
    data_sample = {"words": text.split()}

    test_dataset = UDDataset([data_sample], tokenizer, tags=classes)
    tag_predictions = predict_top_k(model, test_dataset, classes, top_k=top_k_tags)

    test_dataset_l = UDDataset([data_sample], tokenizer_l, tags=classes_l)
    lemma_predictions = predict_top_k(model_l, test_dataset_l, classes_l, top_k=top_k_lemmas)

    result = []
    counter = 1

    for word, tag_options, lemma_options in zip(
        data_sample["words"],
        tag_predictions[0]["top_k_labels"],
        lemma_predictions[0]["top_k_labels"]
    ):

        tag_str = ", ".join([f"{label} ({100*float(prob):.2f}%)" for label, prob in tag_options])
        lemma_str = ", ".join([f"{restore_lemma(word, label)} ({100*float(prob):.2f}%)" for label, prob in lemma_options])
        paradigm_str = ", ".join([f"{label} ({100*float(prob):.2f}%)" for label, prob in lemma_options])

        line = f"{counter}. Form: {word}"
        if show_subtokens == "Yes":
            line += f"\nSubtokens: {' '.join(tokenizer.tokenize(word))}"
        if show_paradigm == "Yes":
            line += f"\nParadigm: {paradigm_str}"
        line += f"\nLemma: {lemma_str}"
        line += f"\nTag: {tag_str}"

        result.append(line)
        result.append("")
        counter += 1

    return "\n".join(result).strip()
  
#--- Интерфейс Gradio ---
demo = gr.Interface(
    fn=analyze_text,
    inputs= [
        gr.Textbox(label="Tokenized sentence", placeholder="Insert tokenized sentence... "),
        gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Top-k for lemmas"),
        gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Top-k for tags"),
        gr.Dropdown(choices=["Yes", "No"],
                    value = "No",
                    label = "Show abstract paradigm label"),
        gr.Dropdown(choices=["Yes", "No"],
                    value = "No",
                    label = "Show subword tokenization"),
    ],
    outputs=gr.Textbox(label="Analysis in UD v2"),
    title="In-context morphological analyzer for Ossetic",
    description="Insert tokenized sentence in Ossetic with spaces around punctuation. Consider prefixes as separate tokens.",
    theme=gr.themes.Base()
)

demo.launch(ssr_mode=False)