Adds documentation and some minor refactors