Few minor bug fixes & cleanup of code.