from utils_dynamic import (
    OUTPUT_DIR,
    ensure_dir,
    plot_smoothed_comparison_from_matrices,
    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 / "figure8_raw")
    max_time_steps = get_max_time_steps_from_env()
    families = [
        (
            "shannon",
            ("Vanilla Shannon Surprise", "Shannon + MISRP", "Shannon + Concept Drift"),
            "Smoothed Mean Error over Time: All Methods",
            "dynamic_error_analysis_all_three_smoothed.pdf",
            "smoothed_error_summary_all_three_methods.csv",
            8100,
        ),
        (
            "bayesian",
            ("Vanilla Bayesian Surprise", "Bayesian + MISRP", "Bayesian + Concept Drift"),
            "Smoothed Mean Error over Time: Bayesian-Based Methods",
            "dynamic_error_analysis_bayesian_all_three_smoothed.pdf",
            "smoothed_error_summary_bayesian_all_three_methods.csv",
            8200,
        ),
        (
            "al1",
            ("Vanilla AL1", "AL1 + MISRP", "AL1 + Concept Drift"),
            "Smoothed Mean Error over Time: AL1-Based Methods",
            "dynamic_error_analysis_al1_all_three_smoothed.pdf",
            "smoothed_error_summary_al1_all_three_methods.csv",
            8300,
        ),
        (
            "al2",
            ("Vanilla AL2", "AL2 + MISRP", "AL2 + Concept Drift"),
            "Smoothed Mean Error over Time: AL2-Based Methods",
            "dynamic_error_analysis_al2_all_three_smoothed.pdf",
            "smoothed_error_summary_al2_all_three_methods.csv",
            8400,
        ),
    ]
    for family, labels, title, plot_file, summary_file, seed in families:
        vanilla = run_vanilla_shannon_experiment(
            seed=seed,
            max_time_steps=max_time_steps,
            data_source="pollution_results",
            gp_normalize_y=True,
        )
        misrp = run_shannon_misrp_experiment(
            seed=seed + 1,
            bin_schedule=False,
            max_time_steps=max_time_steps,
            data_source="pollution_results",
            gp_normalize_y=True,
        )
        concept_drift = run_shannon_misrp_experiment(
            seed=seed + 2,
            bin_schedule=False,
            max_time_steps=max_time_steps,
            data_source="pollution_results",
            gp_normalize_y=False,
        )
        save_generated_experiment(f"figure8_{family}_vanilla", vanilla)
        save_generated_experiment(f"figure8_{family}_misrp", misrp)
        save_generated_experiment(f"figure8_{family}_concept_drift", concept_drift)
        plot_smoothed_comparison_from_matrices(
            labels,
            [vanilla["error_matrix"], misrp["error_matrix"], concept_drift["error_matrix"]],
            title,
            plot_file,
            summary_file,
            out_dir,
        )
    print(f"Saved Figure 8 files from newly generated experiments in {out_dir}")


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