Updated and cleaned code