Clean up code and layout