Final fixes and cleanup