Problem cleaning up the code