removed unnecessary imports
cleaned code a bit