import pandas as pd
import panel as pn
import numpy as np
import hvplot.pandas
from bokeh.models.widgets.tables import NumberFormatter, BooleanFormatter

import matplotlib.pyplot as plt
import seaborn as sns


sns.set(style='whitegrid', context='notebook')
#plt.rcParams["font.family"] = "Latin Modern Sans"


### Data

def load_data(sel_token):
    if sel_token == "virtual":
        df = pd.read_pickle("server_data/virtual_fulltexts_cossim_all_slice=1.pkl")
        df = df.drop("virtual")
    elif sel_token == "boson":
        df = pd.read_pickle("server_data/cossims_plural_slicewidth=1.pkl")
    elif sel_token == "intermediate":
        df = pd.read_pickle("server_data/intermediate_fulltexts_cossim_all_slice=1.pkl")
        df = df.drop("intermediate")

    topn = 1000
    df["sum"] = df.sum(axis=1)
    df = df.astype("float32")
    df = df.round(3)

    df = df.sort_values("sum", ascending=False)
    df = df.head(topn)
    df.index.name = "token"
    return df

sel_token = pn.widgets.Select(
    name="Select dataset",
    value="virtual",
    options=["virtual", "intermediate", "boson"],
    #description="Select the base token ",
)


df = pn.rx(load_data)(sel_token=sel_token)


### Table

table = pn.widgets.Tabulator(
    df,

    ### functionality

    #formatters= {col : NumberFormatter(format='0.000') for col in df.columns}, #tabulator_formatters,
    header_filters = {'token': {'type': 'input', 'func': 'like', 'placeholder': 'search'}},
    selectable='checkbox',

    ### style

    theme = "modern", # 'default', 'site', 'simple', 'midnight', 'modern', 'bootstrap', 'bootstrap4', 'materialize', 'bulma', 'semantic-ui', or 'fast'
    page_size = 8,
    page = 1,
    frozen_columns = {"token" : "left", "sum" : "right"}, # Must give width, otherwise doesn't work!
    width=1800,
    

    ### other

    disabled = True # Whether the cells are editable

)


### Plot

def make_fig():
    fig, ax = plt.subplots(figsize=(12,4))
    df_temp = load_data(sel_token.value)
    if len(table.selection) > 0:
        for i in table.selection:
            df_temp.iloc[i][:-1].plot(ax=ax, label=df_temp.iloc[i].name, lw=2.2, marker=".")
    #else:
        #df.loc["particle"][:-1].plot(ax=ax, label="particle", lw=2.2, marker=".")
        #plt.hist(np.random.random(10))
    plt.ylabel("Cosine Similarity", fontsize=12)
    plt.xlim()
    plt.legend()
    plt.close()

    return fig

def plot_data(event):

    # selected rows as indices in table.selection
    #token = df.iloc[table.selection[0]].name
    #values = df.iloc[table.selection[0]][:-1]
    canvas.loading = True
    fig = make_fig()
    canvas.object = fig
    canvas.loading = False

button = pn.widgets.Button(
    name='Plot', 
    button_type='primary', 
    align="center",
    width=100,
    icon="snowman",
)
button.on_click(plot_data)


canvas = pn.pane.Matplotlib(
    make_fig(),
    format="svg", 
    #width=1000,
    sizing_mode='stretch_width',
    height=400,
    tight=True)


### Serve


ACCENT = "teal"

pn.template.FastListTemplate(
    title="Cosine Similarity for selected tokens",
    sidebar=[],
    main=[pn.Column(
            pn.Row(sel_token),
            table, 
            button, 
            canvas)],
    main_layout=None,
    accent=ACCENT,
).servable()