Updated code to reflect new structure