cleaned up new code a little bit