Updated code organization