Small fixes

* Improve the code readability and apply some better go practices.
* Change way we do sampling: previous is convoluted and error prone.