Adress some comment from the pull request #1076

*Remove some code repetition and fix some error.
*Modifications are based on the comment from the pull request #1076