Updated code with documentation