updated the code a little for cleanliness