better code

Brett told me off. I fixed it.
better user output