Modified code layout to make it more simple