WIP all unit tests passing. TODO refactor and cleanup