Code changes

Updated the code to added missing bits