Clean code

Few commits cleaning up code and getting some TODOs done.

@fcduarte @amitbharadwaj please review