Refactored code as comment on my PR