from utils_dynamic import (
    OUTPUT_DIR,
    ensure_dir,
    plot_shannon_prediction_comparison_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 / "figure11")
    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("figure11_vanilla", vanilla)
    save_generated_experiment("figure11_misrp", misrp)
    plot_shannon_prediction_comparison_from_results(out_dir, misrp, vanilla)
    print(f"Saved Figure 11 files to {out_dir}")


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