Fixed error and cleaned code