Bug fixes & cleanup