Refactor balance implementation to something less stupid