Small changes
Clean up code