added count type and cleaned code