Made adjusments and cleaned code. Fixed #2