Code with done button