"""
---------------------------------------------------------------------------


Air Force Research Laboratory (AFRL) Autonomous Capabilities Team (ACT3)

This is a US Government Work not subject to copyright protection in the US.

The use, dissemination or disclosure of data in this file is subject to
limitation or restriction. See accompanying README and LICENSE for details.
---------------------------------------------------------------------------
Evaluation from Training
"""
import streamlit as st

from corl.evaluation.recording.folder import EpisodeArtifactLoggingCallbackLoader
from corl.visualization.streamlit_app.eval_while_train_layout import StreamlitPageEvaluationWhileTraining

if __name__ == "__main__":
    st.set_page_config(layout="wide")

    # #########################################################################
    # Inspects results generated by running CoRL's separate evaluation script.
    # #########################################################################

    # #########################################################################
    # Directory Selector
    if st.session_state.get("main_page_loaded"):
        st.markdown("# Explore Custom Evaluation Results")
        if not hasattr(st.session_state, "eval_while_train_page"):
            st.session_state.eval_while_train_page = StreamlitPageEvaluationWhileTraining(file_loader=EpisodeArtifactLoggingCallbackLoader)

        st.session_state.eval_while_train_page.display_directory_picker()

        if st.session_state.eval_while_train_page.full_experiment_path:
            st.session_state.eval_while_train_page.display_epoch_picker()
            if st.session_state.eval_while_train_page.epoch_selection != "":  # noqa: PLC1901
                st.session_state.eval_while_train_page.display_load_episode_artifact_option()
            if st.session_state.eval_while_train_page.tables_dict:
                st.session_state.eval_while_train_page.display_episode_metadata_section()
                st.session_state.eval_while_train_page.display_agent_dones_section()
                st.session_state.eval_while_train_page.display_platform_dones_section()
                st.session_state.eval_while_train_page.display_initial_states_section()
                st.session_state.eval_while_train_page.display_units_selector()
                st.session_state.eval_while_train_page.display_agent_steps_section()
                if st.session_state.layout_config:
                    st.session_state.eval_while_train_page.display_trajectory_animation_section(
                        **st.session_state.layout_config["streamlit_layout"]["display_trajectory_animation_section"]
                    )
                else:
                    st.warning("Missing Trajectory animation configuration. No animation is provided.")
                st.session_state.eval_while_train_page.display_space_data_table()
    else:
        st.warning("Main page must be loaded before navigating to subsequent pages. Navigate back to the first page `streamlit_app`)")
