import pandas as pd
import streamlit as st
import plotly.graph_objects as go

def create_sunburst_plot(df):
    fig = go.Figure(go.Sunburst(
        labels=df['labels'],
        parents=df['parents'],
        values=df['values'],
        ids=df['ids'],
        text=df['text'],
        hoverinfo="label+value",
        branchvalues="total",
    ))

    fig.update_layout(margin=dict(t=0, l=0, r=0, b=0))
    return fig

# New data reflecting the CPT codes and expected costs
data = [
    {'ids': 'Root', 'labels': 'Root', 'parents': '', 'values': None, 'text': 'Root'},
    {'ids': 'Hip Surgery', 'labels': 'Hip Surgery', 'parents': 'Root', 'values': 30, 'text': 'Hip Surgery'},
    {'ids': 'Knee Surgery', 'labels': 'Knee Surgery', 'parents': 'Root', 'values': 40, 'text': 'Knee Surgery'},
    {'ids': '99213', 'labels': 'CPT 99213', 'parents': 'Hip Surgery', 'values': 300, 'text': 'Office Visit'},
    {'ids': '99214', 'labels': 'CPT 99214', 'parents': 'Hip Surgery', 'values': 400, 'text': 'Detailed Exam'},
    {'ids': '99284', 'labels': 'CPT 99284', 'parents': 'Knee Surgery', 'values': 250, 'text': 'Emergency Department Visit'},
    {'ids': '99285', 'labels': 'CPT 99285', 'parents': 'Knee Surgery', 'values': 450, 'text': 'Emergency Department Visit, High Complexity'},
]

df = pd.DataFrame(data)

# Filter DataFrame using a query parameter
def filter_data(df, query):
    return df.query(query)

filtered_df = filter_data(df, "parents == 'Root'")

st.plotly_chart(create_sunburst_plot(filtered_df))