More refactoring of states