from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification, TextClassificationPipeline
import torch
import gradio as gr
from openpyxl import load_workbook
from numpy import mean
import pandas as pd
import matplotlib.pyplot as plt

theme = gr.themes.Soft(
    primary_hue="amber",
    secondary_hue="amber",
    neutral_hue="stone",
)

# Load tokenizers and models
tokenizer = AutoTokenizer.from_pretrained("suriya7/bart-finetuned-text-summarization")
model = AutoModelForSeq2SeqLM.from_pretrained("suriya7/bart-finetuned-text-summarization")

tokenizer_keywords = AutoTokenizer.from_pretrained("transformer3/H2-keywordextractor")
model_keywords = AutoModelForSeq2SeqLM.from_pretrained("transformer3/H2-keywordextractor")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
new_model = AutoModelForSequenceClassification.from_pretrained('roberta-rating')
new_tokenizer = AutoTokenizer.from_pretrained('roberta-rating')

classifier = TextClassificationPipeline(model=new_model, tokenizer=new_tokenizer, device=device)

label_mapping = {1: '1/5', 2: '2/5', 3: '3/5', 4: '4/5', 5: '5/5'}

# Function to display and filter the Excel workbook
def filter_xl(file, keywords):
    # Load the workbook and convert it to a DataFrame
    workbook = load_workbook(filename=file)
    sheet = workbook.active
    data = sheet.values
    columns = next(data)[0:]
    df = pd.DataFrame(data, columns=columns)
    
    if keywords:
        keyword_list = keywords.split(',')
        for keyword in keyword_list:
            df = df[df.apply(lambda row: row.astype(str).str.contains(keyword.strip(), case=False).any(), axis=1)]
    
    return df

# Function to calculate overall rating from filtered data
def calculate_rating(filtered_df):
    reviews = filtered_df.to_numpy().flatten()
    ratings = []
    for review in reviews:
        if pd.notna(review):
            rating = int(classifier(review)[0]['label'].split('_')[1])
            ratings.append(rating)
    
    return round(mean(ratings), 2), ratings

# Function to calculate results including summary, keywords, and sentiment
def calculate_results(file, keywords):
    filtered_df = filter_xl(file, keywords)
    overall_rating, ratings = calculate_rating(filtered_df)
    
    # Summarize and extract keywords from the filtered reviews
    text = " ".join(filtered_df.to_numpy().flatten())
    inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors="pt")
    summary_ids = model.generate(inputs["input_ids"], num_beams=2, min_length=10, max_length=50)
    summary = tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    summary = summary.replace("I", "They").replace("my", "their").replace("me", "them")

    inputs_keywords = tokenizer_keywords([text], max_length=1024, truncation=True, return_tensors="pt")
    summary_ids_keywords = model_keywords.generate(inputs_keywords["input_ids"], num_beams=2, min_length=0, max_length=100)
    keywords = tokenizer_keywords.batch_decode(summary_ids_keywords, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

    # Determine overall sentiment
    sentiments = []
    for review in filtered_df.to_numpy().flatten():
        if pd.notna(review):
            sentiment = classifier(review)[0]['label']
            sentiment_label = "Positive" if sentiment == "LABEL_4" or sentiment == "LABEL_5" else "Negative" if sentiment == "LABEL_1" or sentiment == "LABEL_2" else "Neutral"
            sentiments.append(sentiment_label)
    
    overall_sentiment = "Positive" if sentiments.count("Positive") > sentiments.count("Negative") else "Negative" if sentiments.count("Negative") > sentiments.count("Positive") else "Neutral"

    return overall_rating, summary, keywords, overall_sentiment, ratings, sentiments

# Function to analyze a single review
def analyze_review(review):
    if not review.strip():
        return "Error: No text provided", "Error: No text provided", "Error: No text provided", "Error: No text provided"
    
    # Calculate rating
    rating = int(classifier(review)[0]['label'].split('_')[1])
    
    # Summarize review
    inputs = tokenizer([review], max_length=1024, truncation=True, return_tensors="pt")
    summary_ids = model.generate(inputs["input_ids"], num_beams=2, min_length=10, max_length=50)
    summary = tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    summary = summary.replace("I", "he/she").replace("my", "his/her").replace("me", "him/her")

    # Extract keywords
    inputs_keywords = tokenizer_keywords([review], max_length=1024, truncation=True, return_tensors="pt")
    summary_ids_keywords = model_keywords.generate(inputs_keywords["input_ids"], num_beams=2, min_length=0, max_length=100)
    keywords = tokenizer_keywords.batch_decode(summary_ids_keywords, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

    # Determine sentiment
    sentiment = classifier(review)[0]['label']
    sentiment_label = "Positive" if sentiment == "LABEL_4" or sentiment == "LABEL_5" else "Negative" if sentiment == "LABEL_1" or sentiment == "LABEL_2" else "Neutral"

    return rating, summary, keywords, sentiment_label

# Function to count rows in the filtered DataFrame
def count_rows(filtered_df):
    return len(filtered_df)

# Function to plot ratings
def plot_ratings(ratings):
    plt.figure(figsize=(10, 5))
    plt.hist(ratings, bins=range(1, 7), edgecolor='black', align='left')
    plt.xlabel('Rating')
    plt.ylabel('Frequency')
    plt.title('Distribution of Ratings')
    plt.xticks(range(1, 6))
    plt.grid(True)
    plt.savefig('ratings_distribution.png')
    return 'ratings_distribution.png'

# Function to plot sentiments
def plot_sentiments(sentiments):
    sentiment_counts = pd.Series(sentiments).value_counts()
    plt.figure(figsize=(10, 5))
    sentiment_counts.plot(kind='bar', color=['green', 'red', 'blue'])
    plt.xlabel('Sentiment')
    plt.ylabel('Frequency')
    plt.title('Distribution of Sentiments')
    plt.grid(True)
    plt.savefig('sentiments_distribution.png')
    return 'sentiments_distribution.png'

# Gradio interface
with gr.Blocks(theme=theme) as demo:
    gr.Markdown("<h1 style='text-align: center;'>Feedback and Auditing Survey AI Analyzer</h1><br>")
    with gr.Tabs():
        with gr.TabItem("Upload and Filter"):
            with gr.Row():
                with gr.Column(scale=1):
                    excel_file = gr.File(label="Upload Excel File")
                    #excel_file = gr.File(label="Upload Excel File", file_types=[".xlsx", ".xlsm", ".xltx", ".xltm"])
                    keywords_input = gr.Textbox(label="Filter by Keywords (comma-separated)")
                    display_button = gr.Button("Display and Filter Excel Data")
                    clear_button_upload = gr.Button("Clear")
                    row_count = gr.Textbox(label="Number of Rows", interactive=False)
                with gr.Column(scale=3):
                    filtered_data = gr.Dataframe(label="Filtered Excel Contents")
        
        with gr.TabItem("Calculate Results"):
            with gr.Row():
                with gr.Column():
                    overall_rating = gr.Textbox(label="Overall Rating")
                    summary = gr.Textbox(label="Summary")
                    keywords_output = gr.Textbox(label="Keywords")
                    overall_sentiment = gr.Textbox(label="Overall Sentiment")
                    calculate_button = gr.Button("Calculate Results")
                with gr.Column():
                    ratings_graph = gr.Image(label="Ratings Distribution")
                    sentiments_graph = gr.Image(label="Sentiments Distribution")
                    calculate_graph_button = gr.Button("Calculate Graph Results")
        
        with gr.TabItem("Testing Area / Write a Review"):
            with gr.Row():
                with gr.Column(scale=2):
                    review_input = gr.Textbox(label="Write your review here")
                    analyze_button = gr.Button("Analyze Review")
                    clear_button_review = gr.Button("Clear")
                with gr.Column(scale=2):
                    review_rating = gr.Textbox(label="Rating")
                    review_summary = gr.Textbox(label="Summary")
                    review_keywords = gr.Textbox(label="Keywords")
                    review_sentiment = gr.Textbox(label="Sentiment")

    display_button.click(lambda file, keywords: (filter_xl(file, keywords), count_rows(filter_xl(file, keywords))), inputs=[excel_file, keywords_input], outputs=[filtered_data, row_count])
    calculate_graph_button.click(lambda file, keywords: (*calculate_results(file, keywords)[:4], plot_ratings(calculate_results(file, keywords)[4]), plot_sentiments(calculate_results(file, keywords)[5])), inputs=[excel_file, keywords_input], outputs=[overall_rating, summary, keywords_output, overall_sentiment, ratings_graph, sentiments_graph])
    calculate_button.click(lambda file, keywords: (*calculate_results(file, keywords)[:4], plot_ratings(calculate_results(file, keywords)[4])), inputs=[excel_file, keywords_input], outputs=[overall_rating, summary, keywords_output, overall_sentiment])
    analyze_button.click(analyze_review, inputs=review_input, outputs=[review_rating, review_summary, review_keywords, review_sentiment])
    clear_button_upload.click(lambda: (""), outputs=[keywords_input])
    clear_button_review.click(lambda: ("", "", "", "", ""), outputs=[review_input, review_rating, review_summary, review_keywords, review_sentiment])

demo.launch(share=True)