Updated the code structure