from utils_dynamic import (
    OUTPUT_DIR,
    ensure_dir,
    plot_shannon_zoom_from_results,
    get_max_time_steps_from_env,
    run_shannon_misrp_experiment,
    run_vanilla_shannon_experiment,
    save_generated_experiment,
)
from utils_runtime import timed_main


def main():
    out_dir = ensure_dir(OUTPUT_DIR / "figure9_raw")
    max_time_steps = get_max_time_steps_from_env()
    vanilla = run_vanilla_shannon_experiment(seed=16, record_matrix=False, max_time_steps=max_time_steps)
    misrp = run_shannon_misrp_experiment(seed=16, bin_schedule=True, record_matrix=False, max_time_steps=max_time_steps)
    save_generated_experiment("figure9_vanilla", vanilla)
    save_generated_experiment("figure9_misrp", misrp)
    plot_shannon_zoom_from_results(out_dir, misrp, vanilla, view=(0, 300), window_size=20)
    print(f"Saved Figure 9 files to {out_dir}")


if __name__ == "__main__":
    timed_main("figure9_raw.py", main)
