import gradio as gr
import matplotlib.pyplot as plt
import seaborn as sns

from datasets import load_dataset

import histos


dataset = load_dataset("cmpatino/optimal_observables", "train")
dataset_df = dataset["train"].to_pandas()
dataset_df["target"] = dataset_df["target"].map({0: "spin-OFF", 1: "spin-ON"})


def get_plot(features, n_bins):
    plotting_df = dataset_df.copy()
    if len(features) == 1:
        fig, ax = plt.subplots()
        values = [
            plotting_df[plotting_df["target"] == "spin-ON"][features[0]],
            plotting_df[plotting_df["target"] == "spin-OFF"][features[0]],
        ]
        labels = ["spin-ON", "spin-OFF"]
        fig = histos.ratio_hist(
            processes_q=values,
            hist_labels=labels,
            reference_label=labels[1],
            n_bins=n_bins,
            hist_range=None,
            title=features[0],
        )
        return fig
    if len(features) == 2:
        return sns.displot(
            plotting_df,
            x=features[0],
            y=features[1],
            hue="target",
            bins=n_bins,
            height=8,
            aspect=1,
        ).fig


with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            features = gr.Dropdown(
                choices=dataset_df.columns.to_list(),
                label="Feature",
                value="m_tt",
                multiselect=True,
            )
            n_bins = gr.Slider(
                label="Number of Bins for Histogram",
                value=10,
                minimum=10,
                maximum=100,
                step=10,
            )

        feature_plot = gr.Plot(label="Feature's Plot")

    features.change(
        get_plot,
        [features, n_bins],
        feature_plot,
        queue=False,
    )
    n_bins.change(
        get_plot,
        [features, n_bins],
        feature_plot,
        queue=False,
    )
    demo.load(
        get_plot,
        [features, n_bins],
        feature_plot,
        queue=False,
    )

if __name__ == "__main__":
    demo.launch()