# app.py
import streamlit as st
from extractive import preprocess_text, get_sentence_embeddings, build_semantic_graph, apply_textrank, generate_summary
from abstractive import abstractive_summary
from utils import extract_named_entities
from transformers import AutoTokenizer, AutoModel

# Load pre-trained BERT model and tokenizer
model_name = "dmis-lab/biobert-base-cased-v1.2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# Streamlit app layout
st.title("Hybrid Summarization App")
st.write("Upload text files for multi-document summarization or enter text manually for single-document summarization.")

# Multi-document summarization
st.header("Multi-Document Summarization")
uploaded_files = st.file_uploader("Upload text files", type="txt", accept_multiple_files=True)

if uploaded_files:
    texts = [file.read().decode("utf-8") for file in uploaded_files]
    
    # Perform extractive summarization for each document
    extractive_summaries = []
    for text in texts:
        sentences = preprocess_text(text)
        embeddings = get_sentence_embeddings(sentences, model, tokenizer)
        graph = build_semantic_graph(embeddings)
        ranked_sentences = apply_textrank(graph, sentences)
        ext_summary = generate_summary(ranked_sentences, sentences, max_length_ratio=0.5)
        extractive_summaries.append(ext_summary)
    
    # Combine extractive summaries for multi-document summarization
    combined_extractive_summary = " ".join(extractive_summaries)
    st.write("Combined Extractive Summary:", combined_extractive_summary)
    
    # Extract named entities from the combined summary
    entities = extract_named_entities(combined_extractive_summary)
    st.write("Named Entities:", entities)
    
    # Choose summary length ratio for abstractive summarization
    abs_ratio_option = st.selectbox("Choose abstractive summary length ratio", ("1/2", "1/3", "1/4"))
    abs_ratio = {"1/2": 0.5, "1/3": 0.33, "1/4": 0.25}[abs_ratio_option]

    # Perform abstractive summarization
    combined_input = combined_extractive_summary + " " + ' '.join([ent[0] for ent in entities])
    abs_summary = abstractive_summary(combined_input, max_length_ratio=abs_ratio, min_length_ratio=abs_ratio/2)
    st.write("Abstractive Summary:", abs_summary)

# Single-document summarization
st.header("Single-Document Summarization")
text_input = st.text_area("Enter text here")

if text_input:
    # Extract named entities
    entities = extract_named_entities(text_input)
    st.write("Named Entities:", entities)

    # Perform extractive summarization
    sentences = preprocess_text(text_input)
    embeddings = get_sentence_embeddings(sentences, model, tokenizer)
    graph = build_semantic_graph(embeddings)
    ranked_sentences = apply_textrank(graph, sentences)
    ext_summary = generate_summary(ranked_sentences, sentences, max_length_ratio=0.5)
    st.write("Extractive Summary:", ext_summary)

    # Choose summary length ratio for abstractive summarization
    abs_ratio_option = st.selectbox("Choose abstractive summary length ratio", ("1/2", "1/3", "1/4"))
    abs_ratio = {"1/2": 0.5, "1/3": 0.33, "1/4": 0.25}[abs_ratio_option]

    # Perform abstractive summarization
    combined_input = ext_summary + " " + ' '.join([ent[0] for ent in entities])
    abs_summary = abstractive_summary(combined_input, max_length_ratio=abs_ratio, min_length_ratio=abs_ratio/2)
    st.write("Abstractive Summary:", abs_summary)