Updated code logic