Update TODO
Improve modularity and code reuse