Added changes after talking with Yannis