Cleanup/fixes