Apologies for another PR but this stuff is so addictive...

Moved the if statement into the map function for better readability
and added more comments around what I was thinking and why it is the
way it is.