updated and cleaned code