Updated code with some changes.