import streamlit as st
import json

import adrd
ckpt_path = './ckpt_densenet_emb_encoder_2_AUPR.pt'
mdl = adrd.model.ADRDModel.from_ckpt(ckpt_path, device='cpu')

# Create a form for user input
with st.form("json_input_form"):
    st.write("Please enter JSON-formatted dictionary:")
    json_input = st.text_area("JSON Input", height=250)
    submit_button = st.form_submit_button("Predict")

# Process the JSON input
if submit_button:
    try:
        # Parse the JSON input into a Python dictionary
        data_dict = json.loads(json_input)
        _, pred_dict = mdl.predict_proba([data_dict])
        st.write("Predicted probabilities:")
        st.json(pred_dict)
    except json.JSONDecodeError as e:
        # Handle JSON parsing errors
        st.error(f"An error occurred: {e}")