Modified code and readme