More type hints
Minor bug fixes