Updated code structure