import pandas as pd 
import gradio as gr
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt

df = pd.DataFrame({
    "num" : [1, 2 , 3, 4, 5, 6, 7, 8, 9, 10], 
    "channel" : ["alpha", "beta", "beta", "beta", "alpha", "beta", "alpha", "beta", "alpha", "alpha"], 
    "sort_zscore" : [-0.5, -0.3, -0.25, 0.25, 0.4, 0.5, 0.7, 0.75, 0.9, 1.0],
    "D" : [20, 20, 7, 3, 8, 20, 20, 7, 3, 8], 
    "E" : [14, 3, 6, 2, 6, 20, 20, 7, 3, 8]
})

def style_dataframe(df):
    # Compute vmin, vmax from zscore column
    vmin = df['sort_zscore'].min()
    vmax = df['sort_zscore'].max()

    # Create a normalizer and a colormap
    norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
    # cmap = plt.get_cmap('YlGnBu')
    cmap = plt.get_cmap('PuBu')

    def color_row(row):
        """
        For each row:
          1. Get the row's zscore
          2. Normalize it
          3. Convert the normalized value into a hex color from the colormap
          4. Return a list of 'background-color: ...' for each cell
        """
        zscore = row['sort_zscore']
        normalized = norm(zscore)
        color = mcolors.to_hex(cmap(normalized))
        return [f'background-color: {color}'] * len(row)

    # Apply our row-coloring function across every row
    return df.style.apply(color_row, axis=1)

def update_table(channel):
    # Filter based on the selected channel
    filtered_df = df[df["channel"] == channel]
    return style_dataframe(filtered_df)

with gr.Blocks() as demo:
  channel_dropdown = gr.Dropdown(label="Select Channel", choices=["alpha", "beta"], value="beta")
  filtered_table = gr.DataFrame(update_table(channel_dropdown.value), interactive=False, show_search='filter')
  channel_dropdown.change(fn=update_table, inputs=channel_dropdown, outputs=filtered_table)

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