cleaned up code (little bit)