from bokeh.models import ColumnDataSource, FactorRange, HoverTool
from bokeh.plotting import figure
from bokeh.transform import dodge
from bokeh.io import output_notebook, show
from bokeh.palettes import Category20

from bokeh.plotting import figure, output_notebook, show
from bokeh.plotting import figure, show
from bokeh.io import output_file, show
from bokeh.models import ColumnDataSource, FactorRange, Range1d, LinearAxis
from bokeh.transform import factor_cmap
output_notebook()

from huggingface_hub import HfApi
from huggingface_hub import ModelSearchArguments, DatasetSearchArguments
from huggingface_hub import ModelFilter

import pandas as pd
import gradio as gr

api = HfApi()
filt = ModelFilter(library = "diffusers",)
diffusers_models = api.list_models(filter=filt, sort='downloads', direction=-1)
#len(diffusers_models)

diffusers_dict = {}
downloads, authors, modelids, likes = [], [], [], []
print(len(diffusers_models))
for data in diffusers_models:
  #print(data.downloads, data.author, data.modelId, data.likes)
  downloads.append(data.downloads)
  authors.append(data.author) 
  modelids.append(data.modelId) 
  likes.append(data.likes)

diffusers_dict['modelid'] = modelids
diffusers_dict['author'] = authors
diffusers_dict['download'] = downloads
diffusers_dict['likes'] = likes

diffusers_df = pd.DataFrame.from_dict(diffusers_dict)
diffusers_df = diffusers_df[(diffusers_df['download'] != 0) & (diffusers_df['likes'] != 0) ]
grouped = diffusers_df.groupby('author').sum().sort_values(by='download', ascending=False)

#getting data ready for bokeh plots
data_bokeh = grouped.sort_values('download', ascending=False).head(15)
data_bokeh.reset_index(inplace=True)
data_bokeh

#y - axis 1
authors = data_bokeh['author']

#x-axis
downloads = data_bokeh['download']

#y - axis 2
likes = data_bokeh['likes']

# create sample data
data = {'authors': authors, 
        'downloads': downloads, 
        'likes': likes} 


def display_df():
    df = data_bokeh
    return df

def bokehplots():
    source = ColumnDataSource(data=data)
    # set up figure
    p = figure(x_range=FactorRange(*authors), height=350, width=600, title='Downloads and Likes by Author')
    p.vbar(x=dodge('authors',-0.2, range=p.x_range), top='downloads', width=0.4, source=source,
          color=Category20[3][0], legend_label='Downloads')
    p.vbar(x=dodge('authors',0.2, range=p.x_range), top='likes', width=0.4, source=source,
          color=Category20[3][1], legend_label='Likes')
    p.xaxis.major_label_orientation = 45

    # set up y-axis for downloads
    p.yaxis.axis_label = 'Downloads'
    p.yaxis.axis_label_text_color = Category20[3][0]
    p.yaxis.major_label_text_color = Category20[3][0]

    # set up y-axis for likes
    p.extra_y_ranges = {'likes': Range1d(start=0, end=max(likes)+500)}
    p.add_layout(LinearAxis(y_range_name='likes', axis_label='Likes',
                            axis_label_text_color=Category20[3][1], major_label_text_color=Category20[3][1]), 'right')
    p.vbar(x=dodge('authors', 0.4, range=p.x_range), top='likes', width=0.5, source=source,
          color=Category20[3][1], legend_label='Likes', y_range_name='likes')

    # Create a HoverTool object and specify the information to display in the tooltip
    #hover = HoverTool(tooltips=[('authors', '@authors'), ('downloads', '@downloads'), ('likes', '@likes')])
    # Add the HoverTool to the plot
    #p.add_tools(hover)

    # remove grid lines
    p.xgrid.grid_line_color = None
    p.ygrid.grid_line_color = None

    # set legend location
    p.legend.location = 'top_right'

    return p


with gr.Blocks(css = '#myplot {height: 600px;}') as demo: 
    with gr.Row():
        plot = gr.Plot(elem_id='myplot')
        out_dataframe = gr.Dataframe(wrap=True, max_rows=10, overflow_row_behaviour= "paginate", datatype = ["str", "number", "number"], interactive=False)
    demo.load(bokehplots, outputs=[plot])
    demo.load(fn=display_df, outputs=[out_dataframe])  

demo.launch(debug=True)