import streamlit as st
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from html import escape


@st.cache_resource
def load_model_and_tokenizer(path: str):
    with st.spinner("Loading the model... It may take some time"):
        tokenizer = AutoTokenizer.from_pretrained(path)
        model = AutoModelForSequenceClassification.from_pretrained(path, output_attentions=True)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model.to(device)
        model.eval()
    return tokenizer, model, device


def predict_top_p(title, summary, tokenizer, model, device, p=0.95, max_length=320, min_prob=0.01):
    text = f"Title: {title} [SEP] Abstract: {summary}"

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=max_length
    ).to(device)

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.sigmoid(outputs.logits).cpu().squeeze().numpy()

    all_preds = []
    for i, prob_val in enumerate(probs):
        code = model.config.id2label[i]
        all_preds.append((code, float(prob_val)))

    all_preds.sort(key=lambda x: x[1], reverse=True)

    selected = []
    cumulative_prob = 0.0
    for code, prob in all_preds:
        if prob < min_prob:
            break
        selected.append((code, prob))
        cumulative_prob += prob
        if cumulative_prob >= p:
            break

    return selected, cumulative_prob


def predict_top_n(title, summary, tokenizer, model, device, n=5, max_length=320, min_prob=0.01):
    text = f"Title: {title} [SEP] Abstract: {summary}"

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=max_length
    ).to(device)

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.sigmoid(outputs.logits).cpu().squeeze().numpy()

    all_preds = []
    for i, prob_val in enumerate(probs):
        code = model.config.id2label[i]
        all_preds.append((code, float(prob_val)))

    all_preds.sort(key=lambda x: x[1], reverse=True)

    selected = []
    cumulative_prob = 0.0
    for code, prob in all_preds:
        if prob < min_prob:
            continue
        selected.append((code, prob))
        cumulative_prob += prob
        if len(selected) >= n:
            break

    return selected, cumulative_prob


def get_label_index(label, model):
    return int(next(i for i, lbl in model.config.id2label.items() if lbl == label))


def explain_prediction(text, label, tokenizer, model, device, max_length=320):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=max_length
    ).to(device)

    input_ids = inputs["input_ids"]
    attention_mask = inputs["attention_mask"]
    token_type_ids = inputs.get("token_type_ids")
    if token_type_ids is not None:
        token_type_ids = token_type_ids.to(device)

    embeds = model.get_input_embeddings()(input_ids)
    embeds.retain_grad()

    model.zero_grad(set_to_none=True)
    outputs = model(
        inputs_embeds=embeds,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
    )

    label_idx = get_label_index(label, model)
    logit = outputs.logits[0, label_idx]
    prob = torch.sigmoid(logit).item()

    logit.backward()

    grads = embeds.grad[0]
    contrib = (grads * embeds[0]).sum(dim=-1).detach().cpu()

    tokens = tokenizer.convert_ids_to_tokens(input_ids[0].detach().cpu())
    scores = contrib.tolist()

    return tokens, scores, prob


def render_token_heatmap(tokens, scores):
    skip = {"[CLS]", "[SEP]", "[PAD]"}
    items = [(t, s) for t, s in zip(tokens, scores) if t not in skip]
    if not items:
        return "<div>No tokens to display.</div>"

    max_score = max((abs(s) for _, s in items), default=1.0) or 1.0
    parts = []

    for token, score in items:
        norm = abs(score) / max_score
        alpha = 0.12 + 0.88 * norm
        token = token.replace("##", "")
        if score >= 0:
            color = f"rgba(46, 204, 113, {alpha})"
        else:
            color = f"rgba(231, 76, 60, {alpha})"
        parts.append(
            f'<span title="{escape(token)} | {score:+.4f}" '
            f'style="background: {color}; '
            f'padding: 2px 5px; margin: 2px; border-radius: 4px; '
            f'display: inline-block;">{escape(token)}</span>'
        )

    return "<div style='line-height: 2.2;'>" + " ".join(parts) + "</div>"