import streamlit as st


def grid(col=3, row=4, name='grid1'):
    cols = st.columns(col)
    for i in range(row):
        for j in range(col):
            with cols[j]:
                value = st.session_state.checked_dic[name].get(f"{name}_{i*col+j}", False)

                check = st.checkbox(f"{i*col+j}", key=f"{name}_{i*col+j}", value=value)
                if check:
                    st.session_state.checked_dic[name][f"{name}_{i*col+j}"] = True
                else:
                    st.session_state.checked_dic[name][f"{name}_{i*col+j}"] = False


def on_click():
    for key in st.session_state:
        if st.session_state[key] and key[-1].isdigit():
            st.write(key)
    # for key in st.session_state.checked_dic[name]:
    #     if st.session_state.checked_dic[name][key]:
    #         st.write(key)



if __name__ == "__main__":
    if 'checked_dic' not in st.session_state:
        st.session_state.checked_dic = {'grid1': {}, 'grid2': {}}

    name = st.selectbox('Select a grid', ['grid1', 'grid2'])

    with st.form(f"{name}_form"):
        grid(name=name)
        submit_button = st.form_submit_button("Submit", on_click=on_click)