Added code to complete assignment