Clean up code after critique from Kevin