Reworking code and updated TODO.
Reworking code and updated TODO.