Updated code and comments