cleaned my code a bit