updated the code section to be more comprehensive