Cleaned up the state retrieval