Made some changes according to suggestions from @aarondav