from datetime import datetime
from transformers import BartTokenizer, TFBartForConditionalGeneration
from Utils import get_input_chunks
import networkx as nx
from nltk.tokenize import sent_tokenize
from sklearn.feature_extraction.text import TfidfVectorizer
import community


class BARTSummarizer:

    def __init__(self, model_name: str = 'facebook/bart-large-cnn'):
        self.model_name = model_name
        self.tokenizer = BartTokenizer.from_pretrained(model_name)
        self.model = TFBartForConditionalGeneration.from_pretrained(model_name)
        self.max_length = self.model.config.max_position_embeddings

    def summarize(self, text: str):
        encoded_input = self.tokenizer.encode(text, max_length=self.max_length, return_tensors='tf', truncation=True)
        summary_ids = self.model.generate(encoded_input, max_length=300, num_beams=4, early_stopping=True)
        summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        return summary
    
    def chunk_summarize(self, text: str):

        # split the input into chunks
        summaries = []
        input_chunks = get_input_chunks(text, self.max_length)

        # summarize each input chunk separately
        print(datetime.now().strftime("%H:%M:%S"))
        for chunk in input_chunks:
            summaries.append(self.summarize(chunk))
            
        # # combine the summaries to get the final summary for the entire input
        final_summary = " ".join(summaries)

        print(datetime.now().strftime("%H:%M:%S"))

        return final_summary
    
    def preprocess_for_auto_chapters(self, text: str):
       
        # Tokenize the text into sentences
        sentences = sent_tokenize(text)

        # Filter out empty sentences and sentences with less than 5 words
        sentences = [sentence for sentence in sentences if len(sentence.strip()) > 0 and len(sentence.split(" ")) > 4]

        # Combine every 5 sentences into a single sentence
        sentences = [' '.join(sentences[i:i + 5]) for i in range(0, len(sentences), 5)]

        return sentences
    
    def auto_chapters_summarize(self, text: str):

        sentences = self.preprocess_for_auto_chapters(text)

        vectorizer = TfidfVectorizer(stop_words='english')
        X = vectorizer.fit_transform(sentences)

        # Compute the similarity matrix using cosine similarity
        similarity_matrix = X * X.T

        # Convert the similarity matrix to a graph
        graph = nx.from_scipy_sparse_array(similarity_matrix)

        # Apply the Louvain algorithm to identify communities
        partition = community.best_partition(graph, resolution=0.7, random_state=42)

        # Cluster the sentences
        clustered_sentences = []
        for cluster in set(partition.values()):
            sentences_to_print = []
            for i, sentence in enumerate(sentences):
                if partition[i] == cluster:
                    sentences_to_print.append(sentence)
            if len(sentences_to_print) > 1:
                clustered_sentences.append(" ".join(sentences_to_print))
        
        # Summarize each cluster
        summaries = []
        for cluster in clustered_sentences:
            summaries.append(self.chunk_summarize(cluster))

        # Combine the summaries to get the final summary for the entire input
        final_summary = "\n\n".join(summaries)