Improvements as per Tom's comments