from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt

from utils_dynamic import OUTPUT_DIR, ensure_dir, load_pollution_data
from utils_runtime import timed_main


def main():
    out_dir = ensure_dir(OUTPUT_DIR / "figure7")
    normalized_result, _ = load_pollution_data(source="pollution_results")
    grid_size = 50

    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)
    plt.imshow(normalized_result[:, 250].reshape(grid_size, grid_size), cmap="rainbow", interpolation="nearest")
    plt.title("Pollution Map - Time 150")
    plt.colorbar(label="Pollution Level")
    plt.subplot(1, 2, 2)
    plt.imshow(normalized_result[:, 350].reshape(grid_size, grid_size), cmap="rainbow", interpolation="nearest")
    plt.title("Pollution Map - Time 350")
    plt.colorbar(label="Pollution Level")
    plt.tight_layout()
    plt.savefig(out_dir / "pollution_maps.pdf", bbox_inches="tight")
    plt.close()
    print(f"Saved Figure 7 files to {out_dir}")


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