General cleanup and bug fixes