Reviewed code, added/removed some TODOs

- also made some refactoring to enhance readability