Add a loading screen? Cleanup some more code