Cleanup and fixing small bugs.