updated code with proper documentation