Fix run.py
Refactor to a list of states