Cleanup. Minor bug fixes