Using streamlit_folium returned dictionary to actualize the layout of the page

I am trying to make an app that shows a streamlit folium map. Using the dictionary that streamlit_folium returns I want the app to change its layout on the following way:

When clicking on one of the markers on the map would change the layout to double column. The idea is to output information of the marker on the second column.

Unfortunately my current code is not completely responsive and just applies the double column after some clicks. I am not sure why. Can you help me find the error?

As well ideally afterwards i would like to have a button or something that would reinitialise to the original 1 column layout…

Thanks for any help!!

import streamlit as st
import folium
from streamlit_folium import st_folium


st.set_page_config(
        layout="wide",
        initial_sidebar_state="expanded",
    )


# Initialize all necessary session states
if "last_clicked_location" not in st.session_state:
    st.session_state["last_clicked_location"] = None
if "map_data" not in st.session_state:
    st.session_state["map_data"] = None
if "folium_map" not in st.session_state:
    st.session_state["folium_map"] = None

st.title("Folium Map with Multiple Popups & Dynamic Navigation")

# Create map once and store in session state
if st.session_state["folium_map"] is None:
    m = folium.Map(location=[37.7749, -122.4194], zoom_start=5)

    # Add markers
    locations = [
        {"name": "Golden Gate Park", "latitude": 37.7694, "longitude": -122.4862},
        {"name": "Yosemite National Park", "latitude": 37.8651, "longitude": -119.5383},
        {"name": "Grand Canyon", "latitude": 36.1069, "longitude": -112.1129}
    ]

    for loc in locations:
        popup_html = f"""
        <h4>{loc['name']}</h4>
        """
        folium.Marker(
            location=[loc["latitude"], loc["longitude"]],
            popup=folium.Popup(popup_html, max_width=300),
            tooltip=loc["name"],
            icon=folium.Icon(color="green"),
        ).add_to(m)

    st.session_state["folium_map"] = m

# Get cached map
m = st.session_state["folium_map"]

# Handle Click Events
if st.session_state["map_data"] is None:
    # First time initialization
    st.session_state["map_data"] = st_folium(m, key="new")
else:
    # Check if map was clicked
    if isinstance(st.session_state["map_data"], dict):
        clicked_tooltip = st.session_state["map_data"].get("last_object_clicked_tooltip")
        if clicked_tooltip and clicked_tooltip != st.session_state["last_clicked_location"]:
            st.session_state["last_clicked_location"] = clicked_tooltip

# Adjust Layout Dynamically
if st.session_state["last_clicked_location"]:
    left_col, right_col = st.columns([1, 1])
    with left_col:
        st.write("Map View:")
        st.session_state["map_data"] = st_folium(m, width=500, height=600)
    with right_col:
        st.write("### Selected Location Details")
        st.markdown(f"#### {st.session_state['last_clicked_location']}")
else:
    # Show single column with full-width map
    st.session_state["map_data"] = st_folium(m, width=1000, height=600)