Clean up GUI

app.py

@@ -19,22 +19,76 @@ class HvB(HumanVsBot):
             extra_model_environment_context = _extra_model_environment_context,
             **kwargs)
 
-def play(bot_select):
+@st.cache_resource
+def track_play():
+    class trackPlay:
+        def __init__(self):
+            self.active_game = False
+            self.bot_select = None
+        def reset(self):
+            self.active_game = False
+            self.bot_select = None
+    return trackPlay()
+
+def information(play, t_p):
+    display = {
+        True:{True:"Your turn!",False:"Bot is thinking..."},
+        False:{"1-0":"White Wins!","0-1":"Black Wins!","1/2-1/2":"It's a Draw!"},
+    }
+    hum_turn = play.get_human_player() == play.get_curr_player()
+    game_over = play.is_done()
+    choose1 = not game_over
+    choose2 = play.get_result() if game_over else hum_turn
+    return display[choose1][choose2]
+
+def reset(t_p):
+    t_p.reset()
+    st.cache_resource.clear()
+
+def play(bot_select, t_p):
     play = HvB.get_hvb_manager(
         _model=DCMinMax.from_pretrained("mzimm003/{}".format(bot_select)),
         _environment=Chess(render_mode="rgb_array"),
         _extra_model_environment_context=lambda env: {"board":env.board}
     )
-    st.session_state["board"] = streamlit_image_coordinates(play.render_board(), key="brd", click_and_drag=True)
-    play.run()
+    col1, col2 = st.columns([3,1])
+    with col1:
+        st.markdown("# {}".format(information(play, t_p)))
+        st.session_state["board"] = streamlit_image_coordinates(play.render_board(), key="brd", click_and_drag=True)
+    with col2:
+        st.button("Reset", on_click=reset, args=[t_p])
+        st.markdown("#### You are playing as {}!".format("white" if play.get_human_player() == "player_0" else "black"))
+        st.markdown("""
+                    1. Play by dragging the piece you want to move.
+                    2. Illegal moves will not be registered.
+                    """)
+    if not play.is_done():
+        play.run()
 
 def main(kwargs=None):
-    bot_select = st.selectbox(
-        label="bot",
-        options=["DeepChessReplicationMinMax"],
-    )
-    st.button("Play bot!", on_click=play, args=[bot_select])
+    css='''
+    <style>
+        section.main > div {max-width:75rem}
+    </style>
+    '''
+    st.markdown(css, unsafe_allow_html=True)
+
+    t_p = track_play()
+    placeholder = st.empty()
+
+    if not t_p.active_game:
+        # st.cache_resource.clear()
+        with placeholder.container():
+            t_p.bot_select = st.selectbox(
+                label="bot",
+                options=["DeepChessReplicationMinMax"],
+            )
+            t_p.active_game = st.button("Play bot!")
+            if t_p.active_game:
+                st.rerun()
+    else:
+        with placeholder.container():
+            play(t_p.bot_select, t_p)
 
 if __name__ == "__main__":
-    # main()
-    play("DeepChessReplicationMinMax")
+    main()