Refactor code to fix  #84