Minor fixes, added TODO