Code refactor and update of comments