Code cleanup and documentation

[shizuko-akamoto]

• [ ] Clean up TODOs in the code
• [ ] add comments & docs to complex looking functions, edge/special cases