added code to complete assignment