Bug fix(es) and flake8 fixes