import gradio as gr
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
)

### Loading models and tokenizers as global variables

tokenizer_title = AutoTokenizer.from_pretrained("./model_title")
model_title = AutoModelForSequenceClassification.from_pretrained("./model_title")

tokenizer_summary = AutoTokenizer.from_pretrained("./model_summary")
model_summary = AutoModelForSequenceClassification.from_pretrained("./model_summary")

categories = [
    'biology',
    'computer_science',
    'economics',
    'electrical_engineering',
    'finance',
    'mathematics',
    'physics',
    'statistics',
    'unknown'
]

### Auxiliary functions

def classify_title(input_text):
  input_tokens = tokenizer_title.encode(
    input_text,
    truncation=True,
    padding=True,
    return_tensors='pt'
  )

  preds = model_title(input_tokens)
  [probs] = preds.logits.softmax(dim=-1).tolist()
  return {cat: p for (cat, p) in zip(categories, probs)}


def classify_summary(input_text):
  input_tokens = tokenizer_summary.encode(
    input_text,
    truncation=True,
    padding=True,
    return_tensors='pt'
  )

  preds = model_summary(input_tokens)
  [probs] = preds.logits.softmax(dim=-1).tolist()
  return {cat: prob for (cat, prob) in zip(categories, probs)}


def classify_title_summary(title, summary, weight_title=3, weight_summary=1):
  clsf_title = classify_title(title)
  clsf_summary = classify_summary(summary)

  clsf_both = dict()
  for cat in categories:
    clsf_p = clsf_title[cat] * weight_title + clsf_summary[cat] * weight_summary
    clsf_both[cat] = clsf_p / (weight_title + weight_summary)
    
  return clsf_both


def get_top_classifications(clsf, p_top=0.95):    
    clsf_top = dict() 
    p_total = 0

    clsf_sorted = sorted(clsf.items(), key=lambda x: x[1], reverse=True)
    for cat, p in clsf_sorted:
        if p_total >= p_top:
            break
        clsf_top[cat] = p
        p_total += p

    return clsf_top


### Gradio interface 

def process_input(title, summary):
    if not title:
        return None, "Title cannot be empty"
    if len(title.split()) < 3:
        return None, "Title should be at least three words long"

    clsf = classify_title(title) if not summary else classify_title_summary(title, summary)
    clsf = get_top_classifications(clsf, p_top=0.95)
    return clsf, None


title_box = gr.inputs.Textbox(lines=2, placeholder="Your title...", optional=False, label='Article Title')
summary_box = gr.inputs.Textbox(lines=12, placeholder="Your summary...", optional=True, label='Article Summary')
label_widget = gr.outputs.Label(label='Top-95% Predicted Topics')
message_box = gr.outputs.Textbox(label='Errors')

iface = gr.Interface(
    fn=process_input,
    inputs=[title_box, summary_box],
    outputs=[label_widget, message_box],
    title="Article Classifier",
    description="A transformer-based classifier that infers the topic of scientific articles based on a title and summary.",
    allow_flagging="never"
)

iface.launch()