import streamlit as st
import torch
import re
from transformers import AutoTokenizer, AutoModelForCausalLM

HF_MODEL_ID = "dejanseo/latent-entity"


@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID)
    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
    model = AutoModelForCausalLM.from_pretrained(HF_MODEL_ID, torch_dtype=dtype)
    model.eval()
    if torch.cuda.is_available():
        model = model.cuda()
    return tokenizer, model


def strip_md(text):
    text = re.sub(r'\[([^\]]*)\]\([^)]*\)', r'\1', text)
    text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text)
    text = re.sub(r'\*([^*]*)\*', r'\1', text)
    return text


def predict_spans(tokenizer, model, title, text, max_new_tokens=512):
    device = next(model.parameters()).device
    MAX_LEN = 1024
    STRIDE = 256
    COMPLETION_RESERVE = 256

    clean_text = strip_md(text)

    prompt_prefix = f"Title: {title}\n\nText: "
    prompt_suffix = "\n\nHooks:\n"

    prefix_ids = tokenizer(prompt_prefix, add_special_tokens=False)["input_ids"]
    suffix_ids = tokenizer(prompt_suffix, add_special_tokens=False)["input_ids"]
    text_enc = tokenizer(clean_text, add_special_tokens=False)
    text_ids = text_enc["input_ids"]

    fixed_overhead = 1 + len(prefix_ids) + len(suffix_ids)
    text_budget = MAX_LEN - fixed_overhead - COMPLETION_RESERVE

    if text_budget <= 0:
        return []

    all_spans = []
    start = 0

    while start < len(text_ids):
        end = min(start + text_budget, len(text_ids))
        chunk_text_ids = text_ids[start:end]

        input_ids = [tokenizer.bos_token_id] + prefix_ids + chunk_text_ids + suffix_ids
        input_tensor = torch.tensor([input_ids], device=device)

        with torch.no_grad():
            output = model.generate(
                input_tensor,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                eos_token_id=tokenizer.eos_token_id,
            )

        generated_ids = output[0][len(input_ids):]
        generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

        if generated_text and generated_text != "[NONE]":
            for line in generated_text.split("\n"):
                line = line.strip()
                if line and line != "[NONE]":
                    all_spans.append(line)

        if end >= len(text_ids):
            break
        start += STRIDE

    seen = set()
    unique_spans = []
    for span in all_spans:
        if span not in seen:
            seen.add(span)
            unique_spans.append(span)

    return unique_spans


def highlight_spans_in_text(text, spans):
    if not spans:
        return [(text, False)]

    positions = []
    for span in spans:
        idx = text.find(span)
        if idx >= 0:
            positions.append((idx, idx + len(span), span))

    if not positions:
        return [(text, False)]

    positions.sort(key=lambda x: (x[0], -(x[1] - x[0])))
    merged = []
    for s, e, span in positions:
        if merged and s < merged[-1][1]:
            continue
        merged.append((s, e, span))

    segments = []
    pos = 0
    for s, e, span in merged:
        if pos < s:
            segments.append((text[pos:s], False))
        segments.append((text[s:e], True))
        pos = e
    if pos < len(text):
        segments.append((text[pos:], False))

    return segments


st.set_page_config(page_title="Latent Entity Extractor", layout="wide")
st.logo(
    "https://dejan.ai/wp-content/uploads/2024/02/dejan-300x103.png",
    size="large",
    link="https://dejan.ai",
)
st.subheader("Latent Entity Extractor")
st.caption("Identify hidden variables in text — the subjects, reasons, processes, and outcomes that titles withhold to generate clicks. [Read more about latent entities](https://dejan.ai/blog/latent-entities/)")

with st.spinner("Loading model..."):
    tokenizer, model = load_model()

max_tokens = st.slider("Max generation tokens", 64, 1024, 512, 64)

title = st.text_input("Title", placeholder="Enter article title...")
text = st.text_area("Text", height=300, placeholder="Enter article text...")

if st.button("Extract Spans", type="primary") and title and text:
    with st.spinner("Generating..."):
        spans = predict_spans(tokenizer, model, title, text, max_new_tokens=max_tokens)

    if not spans:
        st.warning("No spans predicted.")
    else:
        st.caption(f"{len(spans)} span(s) detected")
        import pandas as pd

        df = pd.DataFrame({"span": spans})
        st.dataframe(df, use_container_width=True, hide_index=True)