Changes in classes to fit the project