import torch
from transformers import XLMRobertaTokenizer
from model.xlmr_for_token_classification import XLMRForTokenClassification
from utils.data_utils import NerProcessor
import gradio as gr

BACKGROUND_COLOR_MAP = {
    'B-PER': '#FFDDDD',
    'I-PER': '#FFAAAA',
    'B-ORG': '#DDDDFF',
    'I-ORG': '#AAAAFF',
    'B-LOC': '#DDFFDD',
    'I-LOC': '#AAFFAA',
    'B-MISC': '#DDDDFF',
    'I-MISC': '#AAAAFF',
    'O': 'white'
}

def load_model(pretrained_path, model_path, device):
    # Load the processor and label map
    processor = NerProcessor()
    label_list = processor.get_labels()
    label_map = {i: label for i, label in enumerate(label_list, 1)}

    # Load the tokenizer
    tokenizer = XLMRobertaTokenizer.from_pretrained(pretrained_path)

    # Initialize the model
    model = XLMRForTokenClassification(
        pretrained_path=pretrained_path,
        n_labels=len(label_list) + 1,
        hidden_size=768,  # or 1024 depending on your model size
        dropout_p=0.1,
        device=device
    )

    # Load the model state
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()

    return model, tokenizer, label_map

def predict_ner(model, tokenizer, text, label_map, device):
    # Tokenize input text
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
    input_ids = inputs["input_ids"].to(device)
    attention_mask = inputs["attention_mask"].to(device)

    # Predict NER labels
    with torch.no_grad():
        logits = model(inputs_ids=input_ids, labels=None, labels_mask=None, valid_mask=None)

    # Convert logits to label IDs
    label_ids = torch.argmax(logits, dim=2).detach().cpu().numpy()

    # Convert label IDs to label names
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
    labels = [label_map[label_id] for label_id in label_ids[0]]

    # Filter out unwanted tokens and remove prefix "▁"
    result = []
    current_word = ''
    current_label = 'O'
    for token, label in zip(tokens, labels):
        if token in ["<s>", "</s>"]:
            continue
        if token.startswith("▁"):
            if current_word:
                result.append((current_word, current_label))
            current_word = token[1:]
            current_label = label
        else:
            current_word += token
        # If the label is not 'O', update the current label to the new label
        if label != 'O':
            current_label = label
    if current_word:
        result.append((current_word, current_label))

    return result

def format_predictions(predictions):
    formatted_text = "<div style='font-size: 22px;'>"  # Increase the font size for the entire text
    for word, label in predictions:
        if label == 'O':
            formatted_text += f'<span>{word}</span> '
        else:
            color = BACKGROUND_COLOR_MAP.get(label, 'white')
            label_style = f'background-color: {color}; font-weight: bold;'
            word_style = f'background-color: {color}; font-size: 22px; padding: 2px;'  # Increase font size for each word
            formatted_text += f'<span style="{word_style}">{word} <span style="{label_style}">[{label}]</span></span> '
    formatted_text += "</div>"
    return formatted_text.strip()

# Load the model and tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, tokenizer, label_map = load_model("./pretrained_models/xlmr.base", "./model.pt", device)

def ner(text):
    predictions = predict_ner(model, tokenizer, text, label_map, device)
    return format_predictions(predictions)

# Define Gradio interface
interface = gr.Interface(
    fn=ner,
    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
    outputs=gr.HTML(),
    title="Named Entity Recognition with XLM-RoBERTa",
    description="This app uses a fine-tuned XLM-RoBERTa model for Named Entity Recognition (NER). Enter a text to see the NER predictions."
)

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