Added a TODO
Updated a comment
Organized imports