import numpy as np
import streamlit as st
from transformers import BartTokenizer, BartForConditionalGeneration
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans
from nltk.tokenize import sent_tokenize
import nltk

# Download required NLTK data
@st.cache_resource
def download_nltk_data():
    try:
        nltk.data.find('tokenizers/punkt')
    except LookupError:
        nltk.download('punkt', quiet=True)
        nltk.download('punkt_tab', quiet=True)

download_nltk_data()

# Load BART model and tokenizer
@st.cache_resource
def load_model():
    tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
    model = BartForConditionalGeneration.from_pretrained('facebook/bart-large-cnn')
    return tokenizer, model

tokenizer, model = load_model()

# Function for single document summarization
def summarize(text, max_length=150, min_length=50):
    inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors='pt')
    summary_ids = model.generate(
        inputs['input_ids'], 
        num_beams=4, 
        max_length=max_length, 
        min_length=min_length, 
        early_stopping=True
    )
    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)

# Function for clustering documents
def cluster_documents(documents, n_clusters=3):
    vectorizer = TfidfVectorizer(stop_words='english')
    X = vectorizer.fit_transform(documents)
    kmeans = KMeans(n_clusters=n_clusters, random_state=42).fit(X)
    return kmeans.labels_

# Function to split long text into chunks
def chunk_text(text, max_words=1000):
    if len(text.split()) <= max_words:
        return [text]
    
    sentences = sent_tokenize(text)
    current_chunk = []
    chunks = []
    current_word_count = 0
    
    for sentence in sentences:
        sentence_words = len(sentence.split())
        if current_word_count + sentence_words <= max_words:
            current_chunk.append(sentence)
            current_word_count += sentence_words
        else:
            chunks.append(" ".join(current_chunk))
            current_chunk = [sentence]
            current_word_count = sentence_words
    
    if current_chunk:
        chunks.append(" ".join(current_chunk))
    
    return chunks

# Function for multi-document summarization
def multi_document_summarize(documents):
    results = {
        'individual_summaries': [],
        'cluster_summaries': [],
        'final_summary': None
    }
    
    # Generate individual document summaries
    for doc in documents:
        chunks = chunk_text(doc)
        doc_summary = " ".join([summarize(chunk) for chunk in chunks])
        results['individual_summaries'].append(doc_summary)
    
    # Cluster documents if we have multiple
    if len(documents) > 1:
        n_clusters = min(3, len(documents))
        clusters = cluster_documents(documents, n_clusters=n_clusters)
        
        # Generate cluster summaries
        for cluster_id in np.unique(clusters):
            cluster_docs = [doc for doc, c in zip(documents, clusters) if c == cluster_id]
            combined_text = " ".join(cluster_docs)
            chunks = chunk_text(combined_text)
            cluster_summary = " ".join([summarize(chunk) for chunk in chunks])
            results['cluster_summaries'].append({
                'doc_indices': [i for i, c in enumerate(clusters) if c == cluster_id],
                'summary': cluster_summary
            })
        
        # Generate final comprehensive summary
        all_summaries = results['individual_summaries'] + [cs['summary'] for cs in results['cluster_summaries']]
        results['final_summary'] = summarize(" ".join(all_summaries), max_length=200, min_length=100)
    else:
        results['final_summary'] = results['individual_summaries'][0]
    
    return results

# Streamlit UI
st.title("Multi-Document Summarization")
st.write("This system summarizes multiple related documents using clustering and BART.")

uploaded_files = st.file_uploader(
    "Upload text files (one document per file)", 
    type=['txt'], 
    accept_multiple_files=True
)

if uploaded_files:
    documents = []
    for file in uploaded_files:
        documents.append(file.read().decode("utf-8"))
    
    if st.button("Generate Summary"):
        with st.spinner("Processing documents..."):
            try:
                results = multi_document_summarize(documents)
                
                # Display individual document summaries
                st.subheader("Individual Document Summaries")
                for i, (doc, summary) in enumerate(zip(documents, results['individual_summaries'])):
                    with st.expander(f"Document {i+1} Summary (First 200 chars: {doc[:200]}...)"):
                        st.write(summary)
                
                # Display cluster summaries if we have multiple documents
                if len(documents) > 1:
                    st.subheader("Cluster Summaries")
                    for i, cluster in enumerate(results['cluster_summaries']):
                        with st.expander(f"Cluster {i+1} Summary (Documents: {', '.join(str(x+1) for x in cluster['doc_indices'])}))"):
                            st.write(cluster['summary'])
                
                # Display final summary
                st.subheader("Final Comprehensive Summary")
                st.write(results['final_summary'])
                
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")
                st.info("Please try again or upload different documents.")
else:
    st.info("Please upload some text files to get started.")