Move code around and add comments