# app.py

from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch
import gradio as gr
import re
from dataclasses import dataclass
from pathlib import Path
import spaces

@dataclass
class SymbolicConfig:
    repo_id: str = "AbstractPhil/bert-beatrix-2048"
    symbolic_roles: list = (
        "<subject>", "<subject1>", "<subject2>", "<pose>", "<emotion>",
        "<surface>", "<lighting>", "<material>", "<accessory>", "<footwear>",
        "<upper_body_clothing>", "<hair_style>", "<hair_length>", "<headwear>",
        "<texture>", "<pattern>", "<grid>", "<zone>", "<offset>",
        "<object_left>", "<object_right>", "<relation>", "<intent>", "<style>",
        "<fabric>", "<jewelry>"
    )

config = SymbolicConfig()
model_dir = snapshot_download(repo_id=config.repo_id)
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForMaskedLM.from_pretrained(model_dir).eval().cuda()


MASK_TOKEN = tokenizer.mask_token or "[MASK]"

@spaces.GPU
def mask_and_predict(text: str, selected_roles: list[str]):
    results = []
    masked_text = text
    token_ids = tokenizer.encode(text, return_tensors="pt").cuda()

    for role in selected_roles:
        role_pattern = re.escape(role)
        masked_text = re.sub(role_pattern, MASK_TOKEN, masked_text)

    masked_ids = tokenizer.encode(masked_text, return_tensors="pt").cuda()
    with torch.no_grad():
        outputs = model(input_ids=masked_ids)
        logits = outputs.logits[0]
        predictions = torch.argmax(logits, dim=-1)

    original_ids = tokenizer.convert_ids_to_tokens(token_ids[0])
    predicted_ids = tokenizer.convert_ids_to_tokens(predictions)
    masked_ids_tokens = tokenizer.convert_ids_to_tokens(masked_ids[0])

    for i, token in enumerate(masked_ids_tokens):
        if token == MASK_TOKEN:
            results.append({
                "Position": i,
                "Masked Token": MASK_TOKEN,
                "Predicted": predicted_ids[i],
                "Original": original_ids[i] if i < len(original_ids) else "",
                "Match": "✅" if predicted_ids[i] == original_ids[i] else "❌"
            })

    accuracy = sum(1 for r in results if r["Match"] == "✅") / max(len(results), 1)
    return results, f"Accuracy: {accuracy:.1%}"

def build_interface():
    role_checkboxes = [gr.Checkbox(label=role, value=False) for role in config.symbolic_roles]

    with gr.Blocks() as demo:
        gr.Markdown("## 🔎 Symbolic BERT Inference Test")
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(label="Symbolic Input Caption", lines=3)
                selected_roles = gr.CheckboxGroup(
                    choices=config.symbolic_roles,
                    label="Mask these symbolic roles"
                )
                run_btn = gr.Button("Run Mask Inference")
            with gr.Column():
                output_table = gr.Dataframe(headers=["Position", "Masked Token", "Predicted", "Original", "Match"], interactive=False)
                accuracy_score = gr.Textbox(label="Mask Accuracy")

        run_btn.click(fn=mask_and_predict, inputs=[input_text, selected_roles], outputs=[output_table, accuracy_score])

    return demo


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