Adjustments after feedback

* Fixed some comments
* fixed coding style to match rest of the code
* better implementation of the actual fix