#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 17 22:46:12 2022
@author: conny
"""

import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

import plotly.express as px

import streamlit as st
from streamlit_option_menu import option_menu

st. set_page_config(layout="wide")

from transformers import pipeline

import pandas as pd

@st.cache(allow_output_mutation = True)
def init_text_summarization_model():
    MODEL = 'facebook/bart-large-cnn'
    pipe = pipeline("summarization", model=MODEL)
    return pipe

@st.cache(allow_output_mutation = True)
def init_zsl_topic_classification():
    MODEL = 'facebook/bart-large-mnli'
    pipe = pipeline("zero-shot-classification", model=MODEL)
    template = "This text is about {}."
    return pipe, template

# Model initialization    
pipeline_summarization = init_text_summarization_model()
pipeline_zsl, template = init_zsl_topic_classification()

st.header('Customer Review Analysis')

# Review text box
default_review = \
    """I attempted to attend the Bank Alpha Eastbank branch last Friday to open a Kids Savings Account for my 3 year old son. I was informed by the Bank Alpha staffer that I was "too close to closing time" and that "I'd have to come back another time". It was about 20 minutes prior to close and there was no one else in the branch except the staff. I did say that I had my son's birth certificate and a Medicare card to establish his identity. Still, no dice. Come back later. No worries, can do.
    I returned today (Monday) to the same branch at 1130, with more than enough time for the account opening to occur. I confirmed with another Bank Alpha staffer that I had my son's birth certificate and Medicare card. However, he went out the back "just to check something". Upon coming back, I was informed that they would not be able to open the account for me today as they required my son, a 3 year old, to be present. The staffer on Friday failed to mention this to me. Equally, the Bank Alpha website for the Kids Savings Account does not list the physical presence of the child as a requirement for opening said account.
    I have never come across a bank so committed to not providing services to prospective customers as Bank Alpha. As a result, my son won't be banking with Bank Alpha, and I probably won't be recommending the use of Bank Alpha to any family or friends either."""
review = st.text_area("Paste/write a review here..", value=default_review, height=250)

tabs = option_menu(menu_title=None,
                   options=[
                       "Text Summarization",
                       "Zero-Shot-Learning",
                       ],
                   default_index=0,
                   orientation='horizontal'
                   )

### Text Summarization
if tabs == 'Text Summarization':
    button = st.button('Summarize review')
    
    if button:
        # Text summarization inference
        with st.spinner("Summarizing review..."):
            summary_text = pipeline_summarization(review, max_length=130, min_length=30, do_sample=False)

        # Show output
        st.write(summary_text[0]['summary_text'])

### Zero-Shot-Learning
elif tabs == 'Zero-Shot-Learning':
    col_product, col_topic = st.columns(2)
    
    # Set product classes
    products = col_product.multiselect(
        label='Available Products and Services:', 
        options=[
            'Bank Account',
            'Credit Card',
            'Home Loan',
            'Insurance',
        ],
        default=[
            'Bank Account',
            'Credit Card',
            'Home Loan',
            'Insurance',
        ]
    )
    product_is_multi_label =  col_product.checkbox("Can have more than one classes", value=True)

    # Set topic classes
    topics = col_topic.multiselect(
        label="Possible Review Topics:",
        options=[
            "Excellent Customer Service", 
            "Great Product Feature", 
            "Poor Service",
            "Unclear Procedure",
            "Other"
        ],
        default=[
            "Excellent Customer Service", 
            "Great Product Feature", 
            "Poor Service",
            "Unclear Procedure",
        ]
    )
    topic_is_multi_label =  col_topic.checkbox("Can have more than one classes", value=False)
    
    button = st.button('Classify')
    
    if button:
        # ZSL inference
        with st.spinner("Identifying product/service and classifying review..."):
            product_classification_output = pipeline_zsl(review, products, hypothesis_template=template, multi_label=product_is_multi_label)
            topic_classification_output = pipeline_zsl(review, topics, hypothesis_template=template, multi_label=topic_is_multi_label)

        # Show output
        col_output_product, col_output_topic = st.columns(2)

        data = {
            'Product': product_classification_output['labels'], 
            'Scores': product_classification_output['scores']
        }
        df = pd.DataFrame(data)
        df = df.sort_values(by='Scores', ascending=True)
        fig = px.bar(df, x='Scores', y='Product', orientation='h')
        col_output_product.plotly_chart(fig, use_container_width=True)

        data = {
            'Topic': topic_classification_output['labels'], 
            'Scores': topic_classification_output['scores']
        }
        df = pd.DataFrame(data)
        df = df.sort_values(by='Scores', ascending=True)
        fig = px.bar(df, x='Scores', y='Topic', orientation='h')
        col_output_topic.plotly_chart(fig, use_container_width=True)