Cleaned up the code. No new functionality