Refactor, fix code formatting and add unittests