Clean this code up a bit