import streamlit as st
from chemical_utils import get_chemical_name, visualize_molecule
from materials_utils import get_mpid, get_data_from_mpid, get_structure_description
import pubchempy as pcp

st.title("Materials and Chemical Data Retrieval")
# st.text("This app retrieves data from the Materials Project API and PubChemPy API.")

if "step" not in st.session_state:
    st.session_state["step"] = 1

with st.sidebar:
    st.header("How to Use the Material Project Retrieval App")
    st.write("""
    This app helps you retrieve data from the Materials Project API and PubChemPy API.
    **Steps:**
    1. Enter a chemical formula and click 'Get MP IDs'.
    2. Select a material from the list.
    3. View detailed information about the selected material.
    """)

if st.session_state["step"] == 1:
    st.header("Enter Chemical Formula")
    chemical = st.text_input("Enter a chemical formula:", "Fe2O3")
    if st.button("Get MP IDs"):
        mpid_dict = get_mpid(chemical)
        if mpid_dict:
            st.session_state["mpid_dict"] = mpid_dict
            st.session_state["chemical"] = chemical
            st.session_state["step"] = 2
        else:
            st.write("No MP IDs found for this formula.")

elif st.session_state["step"] == 2:
    st.header("Select MP ID")
    st.text("Selected Chemical Formula: " + st.session_state["chemical"])
    st.text("Number of MP IDs found: " + str(len(st.session_state["mpid_dict"])))
    st.text("Material IDs(Structure, System, Space Group, Number of Sites, Energy Above Hull, Band Gap):")
    chemical = st.session_state["chemical"]
    mpid_dict = st.session_state["mpid_dict"]

    # Create a display list with both mpid and crystal_system
    display_list = [f"{mpid} ({info[1].value}, {info[2]}, {info[3]}, {info[4]}, {info[5]})" for mpid, info in mpid_dict.items()]
    selected_display = st.selectbox("Select an MP ID:", display_list)

    # Extract the selected mpid
    selected_mpid = selected_display.split(" ")[0]

    if st.button("Get Data"):
        with st.spinner("Retrieving data..."):
            st.write(f"Material ID: {selected_mpid}")
            with st.container(border=True, height=250):
                data_dict = get_data_from_mpid(selected_mpid)
                if data_dict:
                    st.write("Structure Description:")
                    description = get_structure_description(selected_mpid)
                    st.write(description)
                else:
                    st.write("No data found for this material ID.")

    chemical_name = get_chemical_name(chemical)
    if chemical_name:
        st.write(f"Chemical Name: {chemical_name}")
        try:
            compounds = pcp.get_compounds(chemical_name, 'name')
            if compounds:
                smiles = compounds[0].isomeric_smiles
                st.write("3D Structure:(Only in case of Covalent Structure)")
                visualize_molecule(smiles)
            else:
                st.write("No 3D structure found for this formula.")
        except Exception as e:
            st.write(f"Error retrieving 3D structure: {e}")
    else:
        st.write("No chemical name found for this formula.")