Updated code to working version