import os
import torch
import nltk
import gradio as gr
from nltk.tokenize import TweetTokenizer
from transformers import AutoTokenizer, AutoModelForTokenClassification

# 1. NLTK Başlat
nltk.download('punkt', quiet=True)
tknzr = TweetTokenizer(preserve_case=True, strip_handles=False, reduce_len=False)

# 2. Private modele erişim için token al (Space ayarlarından)
hf_token = os.environ.get("HF_TOKEN")

# 3. Modeli Yükle
repo_id = "nypgd/bert-turkish-deprem-ner"
tokenizer = AutoTokenizer.from_pretrained(repo_id, token=hf_token)
model = AutoModelForTokenClassification.from_pretrained(repo_id, token=hf_token)
model.eval()
id2label = model.config.id2label

# 4. Tahmin Fonksiyonu
def extract_entities(text):
    if not text.strip(): return []
    
    original_tokens = tknzr.tokenize(text)
    if not original_tokens: return []

    inputs = tokenizer(original_tokens, is_split_into_words=True, return_tensors="pt", truncation=True)
    with torch.no_grad():
        logits = model(**inputs).logits

    predictions = torch.argmax(logits, dim=2)[0]
    word_ids = inputs.word_ids()

    token_tags = []
    token_logit_idx = []
    prev_word_idx = None
    
    for sub_pos, (pred_id, word_idx) in enumerate(zip(predictions, word_ids)):
        if word_idx is None: continue
        if word_idx != prev_word_idx:
            token_tags.append(id2label[pred_id.item()])
            token_logit_idx.append(sub_pos)
        prev_word_idx = word_idx

    SKIP_TOKENS = {',', '-', ':', '.', '/', '(', ')'}
    entities = []
    current_ent = None

    for i, (token, tag) in enumerate(zip(original_tokens, token_tags)):
        if tag == 'O' and current_ent and token in SKIP_TOKENS: continue
        if tag == 'O':
            if current_ent: entities.append(current_ent)
            current_ent = None
            continue

        ent_type = tag[2:]
        score_val = round(torch.softmax(logits[0, token_logit_idx[i]], dim=-1).max().item(), 4)

        if tag.startswith('B-'):
            if current_ent: entities.append(current_ent)
            current_ent = {"Etiket": f"[{ent_type}]", "Kelime": token, "Skor": score_val}
        elif tag.startswith('I-') and current_ent and current_ent["Etiket"] == f"[{ent_type}]":
            current_ent["Kelime"] += " " + token
            current_ent["Skor"] = round(min(current_ent["Skor"], score_val), 4)
        else:
            if current_ent: entities.append(current_ent)
            current_ent = {"Etiket": f"[{ent_type}]", "Kelime": token, "Skor": score_val}

    if current_ent: entities.append(current_ent)
    
    # Arayüzdeki tablo için veriyi liste formatına çevir
    return [[ent["Etiket"], ent["Kelime"], ent["Skor"]] for ent in entities]

# 5. Gradio Arayüzü Tasarımı
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚨 BERT Türkçe Deprem Tweet NER")
    gr.Markdown("Bu uygulama, deprem tweetlerinden kritik bilgileri (`LOC`, `PER`, `ORG`, `NEED`, `PHONE`, `LINK`) çıkarır. **Model ve Space tamamen gizlidir.**")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(lines=6, label="Tweet Metni", placeholder="Analiz edilecek tweeti buraya yapıştırın...")
            btn = gr.Button("Analiz Et", variant="primary")
            
            gr.Examples(
                examples=[["@AFADTurkiye ve Ahbap ekipleri, Gaziantep İslahiye Yeni Mahalle Karanfil Sokak No:5 adresinde 7 kişi enkaz altında mahsur kaldı. Acil ısıtıcı, çadır ve çocuk maması gerekiyor. Saha sorumlusu Ayşe Yurt iletişim: 0533 123 45 67. Konum ve detaylı bilgi için: https://t.co/yardimadresi"]],
                inputs=input_text,
                label="Örnek Tweet"
            )
            
        with gr.Column(scale=1):
            output_table = gr.Dataframe(headers=["Etiket", "Bulunan Metin", "Güven Skoru"], label="Çıkarılan Varlıklar")
    
    btn.click(fn=extract_entities, inputs=input_text, outputs=output_table)

demo.launch()