import pandas as pd
import streamlit as st
import math


def setup_page():
    st.title("Huggingface model explorer")
    st.text(f"search {len(models_df)} models by name or readme")
    st.text(
        "note that there are many more models but here we only show those with readme"
    )


def setup_inputs():
    col1, col2, col3 = st.columns(3)
    query = col1.text_input("model name query", value="")
    readme_query = col2.text_input("readme query", value="")
    return query, readme_query, col3 


def get_selected_models_df(query, readme_query):
    return models_df[
        models_df["readme"].str.lower().str.contains(readme_query)
        & models_df["modelId"].str.lower().str.contains(query)
    ]


def show_paged_selected_model_info(models_df, page_col):
    page = page_col.number_input("page", 0, math.ceil(len(selected_models_df) / 10))
    selected_models_df_subset = models_df.iloc[
        page * n_per_page : (page+ 1) * n_per_page
    ]
    st.write(f"found {len(selected_models_df)} models")
    for (model_name, tag, readme) in selected_models_df_subset[
        ["modelId", "pipeline_tag", "readme"]
    ].itertuples(index=False):
        with st.expander(f"{model_name} ({tag})"):
            st.write(readme)


models_df = pd.read_csv("models_with_readmes.csv").dropna(subset=["readme"])

setup_page()
query, readme_query, page_col = setup_inputs()
selected_models_df = get_selected_models_df(query, readme_query)

n_per_page = 10

show_paged_selected_model_info(selected_models_df, page_col)