Refactor simulation.py