- improve docstrings
- only pass in the state for a particular input and
not the whole dictionary
- refactor some common code
- minor syntax cleanup