Some fixes and cleanup