Couple of todos added and small tweaks