import wildtorch as wt

wildfire_map = wt.dataset.generate_empty_dataset()

simulator = wt.WildTorchSimulator(
    wildfire_map=wildfire_map,
    simulator_constants=wt.SimulatorConstants(p_continue_burn=0.7),
    initial_ignition=wt.utils.create_ignition(shape=wildfire_map[0].shape),
)

logger = wt.logger.Logger()

for i in range(200):
    simulator.step()
    logger.log_stats(
        step=i,
        num_cells_on_fire=wt.metrics.cell_on_fire(simulator.fire_state).item(),
        num_cells_burned_out=wt.metrics.cell_burned_out(simulator.fire_state).item(),
    )
    logger.snapshot_simulation(simulator)

logger.save_logs()
logger.save_snapshots()