Addresses comments in the GH PR and adds new features.

This commit will be squashed before merging.