added some type hint, still needs more commenting