Minor changes (added TODO)