Small fixes and cleanup