import streamlit as st
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix
from sklearn.preprocessing import LabelEncoder
import pandas as pd
import numpy as np
import time


def get_topic_val(modelname, topics):
    """
    Input: a) modelname: Name of the BERTopic() model used.
           b) topics: The list of topics generated by the model.

    Function: Takes in the input and returns a dictionary with topic names as the keys and the keyword's index values from the df as the values.
    """

    grouped_topics = {topic: [] for topic in set(topics)}

    for index, topic in enumerate(topics):
        grouped_topics[topic].append(index)

    return grouped_topics


def make_result_df(dictionary, topicsdict):
    ''''
    Input: a) Dictionary: The dictionary with results of dict of get_topic_val function. It has all the topics as the keys with their respective keywords as the index given by the model.
           b) topicsdict: The dictionary with the index of the corresponding keyword row in the dataframe as the key and their string keyword as the value.

    Function: Takes in the inputs, maps the index values of keywords with their actual keyword names and returns a result dataframe.
    '''

    key = []
    re_keywords = []
    val = dictionary.items()

    for i, value in val:
        key.append(i)
        val = [*map(topicsdict.get, value)]
        re_keywords.append(val)

    new_dict = {k: v for k, v in zip(key, re_keywords)}
    result_df = pd.DataFrame(new_dict.items(), columns=['Topic Nr', 'Present_Input_Keywords'])
    result_df = result_df.rename_axis('Index').reset_index()

    return result_df


def representativedocs(model, topics, docs, keywords):
    """
    Input: a) Model: Name of the model you want the results for.
           b) topics: topics extracted by the model
           c) docs: documents given as the input to the model. This is the different topic names that the model suggests. (Top n)
           d) Keywords: the input keywords given to the model

    Function: Takes in all the inputs and extracts the representative documents per topic.
    """
    model.get_topic_info()

    # extracting the topic names/numbers
    top_names = model.topic_names
    top_names = pd.DataFrame(top_names.items(), columns=[topics, docs])

    # extracting representative docs for all the topics
    rep_docs = model.representative_docs
    rep_docs = pd.DataFrame(rep_docs.items(), columns=[topics, keywords])

    # get topics with probability
    top_proba = model.get_topics()

    output = pd.merge(top_names,
                      rep_docs,
                      how='left',
                      left_on='topic_num',
                      right_on='topic_num')
    return output


def make_final_dataframe(model, representdocsdf):
    """
    Inputs: a) Model: name of the model
            b) dataframe1: This is the dataframe formed including the topics and their top n topic names for each
            c) representdocsdf: This is the resultant dataframe of the representative docs function


    Function: Returns the resultant dataframe with topic number, their top n names with c-tf-idf scores and all the keywords they contain.
    """
    dataframe1 = pd.DataFrame(model.topics.items(), columns=['Topic Nr', 'Possible Topic Names'])
    finaldfname = pd.merge(dataframe1, representdocsdf)

    return finaldfname


def get_similarity_score(model, topicnr, resultdf, threshold):
    '''
    Parameters:
        Inputs: a) model: the model used to train your topic modelling
                b) topicnr: the topic for which you want to see the similarity score. IMP: here the nr is the index of the row and not the topic nr so for topic -1 = topicnr is 0
                c) resultdf: the resultant df to merge to get combined results
                d) threshold: the threshold above which you want to get similar topics
        Ouput: A pandas dataframe with topicnr, topic names, keywords present (input) and the distance score for each.
    '''

    if model.topic_embeddings is not None:
        embeddings = np.array(model.topic_embeddings)
    else:
        embeddings = model.c_tf_idf

    distance_matrix = cosine_similarity(embeddings)
    data = distance_matrix[topicnr]
    score_df = pd.DataFrame(data=data, columns={'similarity_score'})
    score_df = score_df.rename_axis('Index').reset_index()

    # merging score with resultant dataframe
    df = pd.merge(score_df, resultdf, on='Index')

    df = df[df['similarity_score'] >= threshold]

    return df


def get_similarity_score(model):
    '''
    Parameters:
        Inputs: a) model: the model used to train your topic modelling
                b) topicnr: the topic for which you want to see the similarity score. IMP: here the nr is the index of the row and not the topic nr so for topic -1 = topicnr is 0
                c) resultdf: the resultant df to merge to get combined results
                d) threshold: the threshold above which you want to get similar topics
        Ouput: A pandas dataframe with topicnr, topic names, keywords present (input) and the distance score for each.
    '''

    topics = sorted(list(model.get_topics().keys()))

    # Extract topic words and their frequencies
    topic_list = sorted(topics)

    embeddings = model.c_tf_idf
    distance_matrix = cosine_similarity(embeddings)

    most_similar_ind = []
    most_similar_val = []

    for topic in topic_list:
        data = distance_matrix[topic]  # topic -1
        i = np.argsort(data, axis=0)[-2]
        most_similar_ind.append(i)  # ensure length and order for the list
        most_similar_val.append(data[i])

    similar_df = pd.DataFrame()
    similar_df['Topic Nr'] = topic_list
    similar_df['most_similar'] = most_similar_ind
    similar_df['similarity_score'] = most_similar_val
    return similar_df


def plot_confusion_matrix(model, true, predicted, xfig, yfig, title, topic):
    fig, ax = plt.subplots(figsize=(xfig, yfig))
    ax = sns.heatmap(confusion_matrix(predicted, true),
                     annot=True,
                     fmt='d',
                     cmap="PiYG",
                     xticklabels = topic,
                     yticklabels=topic)
    plt.ylabel('True Values')
    plt.xlabel('Predicted Values')
    plt.title(title)
    plt.xticks(rotation=90, ha='right')
    plt.yticks(rotation=360, ha='right')
    plt.show()
    st.pyplot(fig)



def _update_progress(text, progress):
    text = f"{text}"
    st.write(text, unsafe_allow_html=True)
    progress_bar = st.progress(0)
    progress_bar.progress(progress)

@st.cache
def convert_df(df):
     # IMPORTANT: Cache the conversion to prevent computation on every rerun
    return df.to_csv().encode('utf-8')


def return_top_n_pred_prob_df(n, model, docs, column_name):
    '''
    Function to predict the top n topics for a specific keyword with it's accuracy score
    Parameters:
      Input:
        a) n = Top n topic classes you want
        b) model = the model you have trained your dataset on
        c) docs = the keywords on which you want to predict the top n topics
        d) column_name = name of the column in the resultant df which takes in this as it's input for naming it

      Output: A dataframe with keywords and their corresponding topic names with its associated percentage accuracy.
    '''
    predictions = model.predict_proba(docs)
    preds_idx = np.argsort(-predictions, axis=1)
    top_n_preds = pd.DataFrame()

    for i in range(n):
        top_n_preds['keywords'] = docs
        top_n_preds[column_name + "_" + '{}'.format(i)] = [preds_idx[doc][i] for doc in range(len(docs))]
        top_n_preds[column_name + "_" + '{}_prob'.format(i)] = [predictions[doc][preds_idx[doc][i]] for doc in
                                                                range(len(docs))]

        top_n_preds = top_n_preds.rename(columns={'class_name': column_name + ''.format(i)})
        try:
            top_n_preds.drop(columns=['index', column_name + '_prediction_{}_num'.format(i)], inplace=True)
        except:
            pass
    return top_n_preds


from sklearn.metrics.pairwise import cosine_similarity


def get_class_similarity_score(model):
    '''
    Parameters:
        Inputs: a) model: the model used to train your topic modelling
                b) topicnr: the topic for which you want to see the similarity score. IMP: here the nr is the index of the row and not the topic nr so for topic -1 = topicnr is 0
                c) resultdf: the resultant df to merge to get combined results
                d) threshold: the threshold above which you want to get similar topics
        Ouput: A pandas dataframe with topicnr, topic names, keywords present (input) and the distance score for each.
    '''

    topics = sorted(list(model.get_topics().keys()))

    # Extract topic words and their frequencies
    topic_list = sorted(topics)

    embeddings = model.c_tf_idf
    distance_matrix = cosine_similarity(embeddings)

    most_similar_ind = []
    most_similar_val = []

    for topic in topic_list:
        data = distance_matrix[topic]  # topic -1
        i = np.argsort(data, axis=0)[-2]
        most_similar_ind.append(i)  # ensure length and order for the list
        most_similar_val.append(data[i])

    similar_df = pd.DataFrame()
    similar_df['Topic Nr'] = topic_list
    similar_df['most_similar'] = most_similar_ind
    similar_df['c_similarity_score'] = most_similar_val
    similar_df = similar_df.rename_axis('Index').reset_index()
    return similar_df