simplify the code a bit. Add comments for what I've worked out