Updated Source Code to complete assignment