# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging

import gradio as gr
import numpy as np
import pandas as pd
from matplotlib.figure import Figure
import seaborn as sns
import statistics
# TODO: Temporary (for WiFi reasons)
# import streamlit as st
import utils
import utils.dataset_utils as ds_utils
# from st_aggrid import AgGrid, GridOptionsBuilder
from utils.dataset_utils import HF_DESC_FIELD, HF_FEATURE_FIELD, HF_LABEL_FIELD

logs = utils.prepare_logging(__file__)
# TODO: Toggled off for WiFi reasons
# st.set_option('deprecation.showPyplotGlobalUse', False)

# Note: Make sure to consider colorblind-friendly colors for your images! Ex:
# ["#332288", "#117733", "#882255", "#AA4499", "#CC6677", "#44AA99", "#DDCC77",
# "#88CCEE"]

pd.options.display.float_format = "{:,.3f}".format # '{:20,.2f}'.format

def subheader():
    gr.Markdown("""This demo showcases the 
    [dataset metrics as we develop them](https://huggingface.co/blog/data-measurements-tool).
    Right now this has:
    - dynamic loading of datasets in the lib
    - fetching config and info without downloading the dataset
    - propose the list of candidate text and label features to select.
    """)


def get_label_names(dataset_name: str, config_name: str, ds_name_to_dict):
    label_field, label_names = (
        ds_name_to_dict[dataset_name][config_name][HF_FEATURE_FIELD][
            HF_LABEL_FIELD][0]
        if len(
            ds_name_to_dict[dataset_name][config_name][HF_FEATURE_FIELD][
                HF_LABEL_FIELD]
        ) > 0
        else ((), [])
    )
    return label_field, label_names



def update_dataset(dataset_name: str, ds_name_to_dict):
    # choose a config to analyze
    ds_configs = ds_name_to_dict[dataset_name]
    # special handling for the largest-by-far dataset, C4
    if dataset_name == "c4":
        config_names = ['en', 'en.noblocklist', 'realnewslike']
    else:
        config_names = list(ds_configs.keys())

    config_name = config_names[0]
    ds_config = ds_configs[config_name]

    text_features = ds_config[HF_FEATURE_FIELD]["string"]
    text_features = [('text',)] if dataset_name == "c4" else [tp for tp in text_features if tp[0] != "id"]
    feature = str(text_features[0])
    text_features = [str(f) for f in text_features]

    avail_splits = list(ds_config["splits"].keys())
    split = avail_splits[0]

    return [(config_names, config_name), (text_features, feature), (avail_splits, split)]


def update_config(dataset_name: str, config_name: str, ds_name_to_dict):
    ds_config = ds_name_to_dict[dataset_name][config_name]

    text_features = ds_config[HF_FEATURE_FIELD]["string"]
    text_features = [('text',)] if dataset_name == "c4" else [tp for tp in text_features if tp[0] != "id"]
    feature = str(text_features[0])
    text_features = [str(f) for f in text_features]

    avail_splits = list(ds_config["splits"].keys())
    split = avail_splits[0]

    return [(text_features, feature), (avail_splits, split)]


def sidebar_selection(ds_name_to_dict, column_id=""):
    ds_names = list(ds_name_to_dict.keys())
    with gr.Accordion(f"Choose dataset and field {column_id}", open=True):
        subheader()
        # choose a dataset to analyze
        ds_name = gr.Dropdown(
            label=f"Choose dataset to explore{column_id}:",
            choices=ds_names,
            value="hate_speech18",
        )
        # choose a config to analyze
        ds_configs = ds_name_to_dict[ds_name.value]
        # special handling for the largest-by-far dataset, C4
        if ds_name == "c4":
            config_names = ['en', 'en.noblocklist', 'realnewslike']
        else:
            config_names = list(ds_configs.keys())
        config_name = gr.Dropdown(
            label=f"Choose configuration{column_id}:",
            choices=config_names,
            value=config_names[0],
        )
        # choose a subset of num_examples
        ds_config = ds_configs[config_name.value]
        text_features = ds_config[HF_FEATURE_FIELD]["string"]
        # TODO @yacine: Explain what this is doing and why eg tp[0] could = "id"
        text = f"Which text feature from the {column_id} dataset would you like to analyze?"
        choices = [('text',)] if ds_name == "c4" else [tp for tp in text_features if tp[0] != "id"]
        text_field = gr.Dropdown(
            label=text,
            choices=[str(f) for f in choices],
            value=str(choices[0])
        )
        # Choose a split and dataset size
        avail_splits = list(ds_config["splits"].keys())
        # 12.Nov note: Removing "test" because those should not be examined
        # without discussion of pros and cons, which we haven't done yet.
        if "test" in avail_splits:
            avail_splits.remove("test")
        split = gr.Dropdown(
            label=f"Which split from the{column_id} dataset would you like to analyze?",
            choices=avail_splits,
            value=avail_splits[0],
        )
        label_field, label_names = get_label_names(ds_name.value, config_name.value, ds_name_to_dict)
        calculate_btn = gr.Button(value="Calculate", variant="primary")
        return {
            "dset_name": ds_name,
            "dset_config": config_name,
            "split_name": split,
            "text_field": text_field,
            "label_field": label_field,
            "label_names": label_names,
            "calculate_btn": calculate_btn
        }


def expander_header(dstats, ds_name_to_dict, column_id=""):
    with st.expander(f"Dataset Description{column_id}"):
        st.markdown(
            ds_name_to_dict[dstats.dset_name][dstats.dset_config][HF_DESC_FIELD]
        )
        st.dataframe(dstats.dset_peek)


def expander_general_stats(dstats, column_id=""):
    with gr.Accordion(f"General Text Statistics{column_id}"):
        st.caption(
            "Use this widget to check whether the terms you see most "
            "represented in the dataset make sense for the goals of the dataset."
        )
        st.markdown("There are {0} total words".format(str(dstats.total_words)))
        st.markdown(
            "There are {0} words after removing closed "
            "class words".format(str(dstats.total_open_words))
        )
        st.markdown(
            "The most common "
            "[open class words](https://dictionary.apa.org/open-class-words) "
            "and their counts are: "
        )
        st.dataframe(dstats.sorted_top_vocab_df)
        st.markdown(
            "There are {0} missing values in the dataset.".format(
                str(dstats.text_nan_count)
            )
        )
        if dstats.dups_frac > 0:
            st.markdown(
                "The dataset is {0}% duplicates. "
                "For more information about the duplicates, "
                "click the 'Duplicates' tab below.".format(
                    str(round(dstats.dups_frac * 100, 2)))
            )
        else:
            st.markdown("There are 0 duplicate items in the dataset. ")


def expander_label_distribution(dstats, column_id=""):
    with st.expander(f"Label Distribution{column_id}", expanded=False):
        st.caption(
            "Use this widget to see how balanced the labels in your dataset are."
        )
        if dstats.fig_labels:
            st.plotly_chart(dstats.fig_labels, use_container_width=True)
        else:
            st.markdown("No labels were found in the dataset")


def expander_text_lengths(dstats, column_id=""):
    _TEXT_LENGTH_CAPTION = (
        "Use this widget to identify outliers, particularly suspiciously long "
        "outliers."
    )
    with st.expander(f"Text Lengths{column_id}", expanded=False):
        st.caption(_TEXT_LENGTH_CAPTION)
        st.markdown(
            "Below, you can see how the lengths of the text instances in your "
            "dataset are distributed."
        )
        st.markdown(
            "Any unexpected peaks or valleys in the distribution may help to "
            "identify instances you want to remove or augment."
        )
        st.markdown(
            "### Here is the count of different text lengths in "
            "your dataset:"
        )
        # When matplotlib first creates this, it's a Figure.
        # Once it's saved, then read back in,
        # it's an ndarray that must be displayed using st.image
        # (I know, lame).
        if isinstance(dstats.length_obj.fig_lengths, Figure):
            st.pyplot(dstats.length_obj.fig_lengths, use_container_width=True)
        else:
            try:
                st.image(dstats.length_obj.fig_lengths)
            except Exception as e:
                logs.exception("Hit exception for lengths figure:")
                logs.exception(e)
        st.markdown(
            "The average length of text instances is **"
            + str(round(dstats.length_obj.avg_length, 2))
            + " words**, with a standard deviation of **"
            + str(round(dstats.length_obj.std_length, 2))
            + "**."
        )
        if dstats.length_obj.lengths_df is not None:
            start_id_show_lengths = st.selectbox(
                "Show examples of length:",
                np.sort(dstats.length_obj.lengths_df["length"].unique())[::-1].tolist(),
                key=f"select_show_length_{column_id}",
            )
            st.table(
                dstats.length_obj.lengths_df[
                    dstats.length_obj.lengths_df["length"] == start_id_show_lengths
                ].set_index("length")
            )


def expander_text_duplicates(dstats, column_id=""):
    with st.expander(f"Text Duplicates{column_id}", expanded=False):
        st.caption(
            "Use this widget to identify text strings that appear more than "
            "once."
        )
        st.markdown(
            "A model's training and testing may be negatively affected by "
            "unwarranted duplicates "
            "([Lee et al., 2021](https://arxiv.org/abs/2107.06499))."
        )
        st.markdown("------")
        st.write(
            "### Here is the list of all the duplicated items and their counts "
            "in the dataset."
        )
        if not dstats.duplicates_results:
            st.write("There are no duplicates in this dataset! 🥳")
        else:
            st.write("The fraction of the data that is a duplicate is:")
            st.write(str(round(dstats.dups_frac, 4)))
            # TODO: Check if this is slow when the size is large --
            # Should we store as dataframes?
            # Dataframes allow this to be interactive.
            st.dataframe(ds_utils.counter_dict_to_df(dstats.dups_dict))


def expander_text_perplexities(dstats, column_id=""):
    with st.expander(f"Text Perplexities{column_id}", expanded=False):
        st.caption(
            "Use this widget to identify text perplexities from GPT-2."
        )
        st.markdown(
            """
            Outlier perplexities, especially very high values, could highlight 
            an issue with an example. Smaller variations should be interpreted 
            with more care, as they indicate how similar to the GPT-2 training 
            corpus the examples are rather than being reflective of general 
            linguistic properties.
            For more information on GPT-2, 
            see its [model card](https://hf.co/gpt2).
            """
        )
        st.markdown("------")
        st.write(
            "### Here is the list of the examples in the dataset, sorted by "
            "GPT-2 perplexity:"
        )
        if dstats.perplexities_df is None or dstats.perplexities_df.empty:
            st.write(
                "Perplexities have not been computed yet for this dataset, or "
                "this dataset is too large for the UI (> 1,000,000 examples).")
        else:
            st.dataframe(dstats.perplexities_df.reset_index(drop=True))


def expander_npmi_description(min_vocab):
    _NPMI_CAPTION = (
        "Use this widget to identify problematic biases and stereotypes in "
        "your data."
    )
    _NPMI_CAPTION1 = """
    nPMI scores for a word help to identify potentially
    problematic associations, ranked by how close the association is."""
    _NPMI_CAPTION2 = """
    nPMI bias scores for paired words help to identify how word
    associations are skewed between the selected selected words
    ([Aka et al., 2021](https://arxiv.org/abs/2103.03417)).
    """

    st.caption(_NPMI_CAPTION)
    st.markdown(_NPMI_CAPTION1)
    st.markdown(_NPMI_CAPTION2)
    st.markdown("  ")
    st.markdown(
        "You can select from gender and sexual orientation "
        "identity terms that appear in the dataset at least %s "
        "times." % min_vocab
    )
    st.markdown(
        "The resulting ranked words are those that co-occur with both "
        "identity terms.  "
    )
    st.markdown(
        "The more *positive* the score, the more associated the word is with "
        "the first identity term.  "
        "The more *negative* the score, the more associated the word is with "
        "the second identity term."
    )


def expander_zipf(dstats, column_id=""):
    z = dstats.z
    zipf_fig = dstats.zipf_fig
    with st.expander(
        f"Vocabulary Distribution{column_id}: Zipf's Law Fit", expanded=False
    ):
        try:
            _ZIPF_CAPTION = """This shows how close the observed language is to an ideal
            natural language distribution following [Zipf's law](https://en.wikipedia.org/wiki/Zipf%27s_law),
            calculated by minimizing the [Kolmogorov-Smirnov (KS) statistic](https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Smirnov_test)."""

            powerlaw_eq = r"""p(x) \propto x^{- \alpha}"""
            zipf_summary = (
                    "The optimal alpha based on this dataset is: **"
                    + str(round(z.alpha, 2))
                    + "**, with a KS distance of: **"
                    + str(round(z.ks_distance, 2))
            )
            zipf_summary += (
                    "**.  This was fit with a minimum rank value of: **"
                    + str(int(z.xmin))
                    + "**, which is the optimal rank *beyond which* the scaling regime of the power law fits best."
            )

            alpha_warning = "Your alpha value is a bit on the high side, which means that the distribution over words in this dataset is a bit unnatural. This could be due to non-language items throughout the dataset."
            xmin_warning = "The minimum rank for this fit is a bit on the high side, which means that the frequencies of your most common words aren't distributed as would be expected by Zipf's law."
            fit_results_table = pd.DataFrame.from_dict(
                {
                    r"Alpha:": [str("%.2f" % z.alpha)],
                    "KS distance:": [str("%.2f" % z.ks_distance)],
                    "Min rank:": [str("%s" % int(z.xmin))],
                },
                columns=["Results"],
                orient="index",
            )
            fit_results_table.index.name = column_id
            st.caption(
                "Use this widget for the counts of different words in your dataset, measuring the difference between the observed count and the expected count under Zipf's law."
            )
            st.markdown(_ZIPF_CAPTION)
            st.write(
                """
            A Zipfian distribution follows the power law: $p(x) \propto x^{-α}$
    with an ideal α value of 1."""
            )
            st.markdown(
                "In general, an alpha greater than 2 or a minimum rank greater than 10 (take with a grain of salt) means that your distribution is relativaly _unnatural_ for natural language. This can be a sign of mixed artefacts in the dataset, such as HTML markup."
            )
            st.markdown(
                "Below, you can see the counts of each word in your dataset vs. the expected number of counts following a Zipfian distribution."
            )
            st.markdown("-----")
            st.write("### Here is your dataset's Zipf results:")
            st.dataframe(fit_results_table)
            st.write(zipf_summary)
            # TODO: Nice UI version of the content in the comments.
            # st.markdown("\nThe KS test p-value is < %.2f" % z.ks_test.pvalue)
            # if z.ks_test.pvalue < 0.01:
            #    st.markdown(
            #        "\n Great news! Your data fits a powerlaw with a minimum KS " "distance of %.4f" % z.distance)
            # else:
            #    st.markdown("\n Sadly, your data does not fit a powerlaw. =(")
            # st.markdown("Checking the goodness of fit of our observed distribution")
            # st.markdown("to the hypothesized power law distribution")
            # st.markdown("using a Kolmogorov–Smirnov (KS) test.")
            st.plotly_chart(zipf_fig, use_container_width=True)
            if z.alpha > 2:
                st.markdown(alpha_warning)
            if z.xmin > 5:
                st.markdown(xmin_warning)
        except:
            st.write("Under construction!")


def npmi_widget(dstats, column_id=""):
    """
    Part of the UI, but providing for interaction.
    :param column_id:
    :param dstats:
    :return:
    """
    min_vocab = dstats.min_vocab_count
    npmi_stats = dstats.npmi_obj
    available_terms = npmi_stats.avail_identity_terms
    with st.expander(f"Word Association{column_id}: nPMI", expanded=False):
        if npmi_stats and len(available_terms) > 0:
            expander_npmi_description(min_vocab)
            st.markdown("-----")
            term1 = st.selectbox(
                f"What is the first term you want to select?{column_id}",
                available_terms,
            )
            term2 = st.selectbox(
                f"What is the second term you want to select?{column_id}",
                reversed(available_terms),
            )
            try:
                joint_npmi_df = npmi_stats.get_display(term1, term2)
                npmi_show(joint_npmi_df)
            except Exception as e:
                logs.exception(e)
                st.markdown(
                    "**WARNING!** The nPMI for these terms has not been"
                    " pre-computed, please re-run caching."
                )
        else:
            st.markdown("No words found co-occurring with both of the selected identity"
                " terms.")


def npmi_show(paired_results):
    if paired_results.empty:
        st.markdown(
            "No words that co-occur enough times for results! Or there's a 🐛."
            "  Or we're still computing this one. 🤷")
    else:
        logs.debug("Results to be shown in streamlit are")
        logs.debug(paired_results)
        s = pd.DataFrame(
            paired_results.sort_values(paired_results.columns[0], ascending=True))
        s.index.name = "word"
        bias_col = s.filter(like="bias").columns
        #count_cols = s.filter(like="count").columns
        # Keep the dataframe from being crazy big.
        if s.shape[0] > 10000:
            bias_thres = max(abs(s[s[0]][5000]),
                             abs(s[s[0]][-5000]))
            logs.info(f"filtering with bias threshold: {bias_thres}")
            s_filtered = s[s[0].abs() > bias_thres]
        else:
            s_filtered = s
        cm = sns.palplot(sns.diverging_palette(270, 36, s=99, l=48, n=16))
        out_df = s_filtered.style.background_gradient(subset=bias_col, cmap=cm).format(formatter="{:,.3f}").set_properties(**{"align": "center", "width":"100em"}).set_caption("nPMI scores between the selected identity terms and the words they both co-occur with")
        #set_properties(subset=count_cols, **{"width": "10em", "text-align": "center"}).
        # .format(subset=count_cols, formatter=int).
        #.format(subset=bias_col, formatter="{:,.3f}")
        st.write("### Here is your dataset's bias results:")
        st.dataframe(out_df)
