More bug fixes :(