Fixed small bug. Updated TODO.
Small changes.