Clean up

Made it work a little better.