Code improvement and bug fixes