import pickle as pkl
import plotly.express as px
import gradio as gr

# Load the job stack count data
with open("job_stack_count.pkl", "rb") as f:
    job_stack_count = pkl.load(f)

# Load the user's stack data
with open("all_my_stacks.pkl", "rb") as f:
    mynew_stacks = pkl.load(f)

def generate_treemap(category: str):
    """Generate a treemap visualization for the selected category."""
    if category not in job_stack_count:
        return "Category not found"
    
    stat_dict = job_stack_count[category]
    stat_dict_keys = list(stat_dict.keys())
    stat_dict_values = list(stat_dict.values())
    total = len(mynew_stacks)
    
    stat_dict_perc = [val / total * 100 for val in stat_dict_values]
    hover_text = [f"{label}: Covers {p:.0f}% of Job Descriptions" for label, p in zip(stat_dict_keys, stat_dict_perc)]
    
    fig = px.treemap(
        names=stat_dict_keys,
        parents=[""] * len(stat_dict_keys),  # No parent categories
        values=stat_dict_values,
        hover_name=hover_text
    )
    fig.update_layout(title=category)
    return fig

# Create a Gradio interface
categories = list(job_stack_count.keys())
demo = gr.Interface(
    fn=generate_treemap,
    inputs=gr.Dropdown(choices=categories, label="Select Category"),
    outputs=gr.Plot(),
    title="Tech Stack Treemap Visualization",
    description="Select a category to visualize the distribution of job tech stacks."
)

# Launch the app
demo.launch(share=True)