clean up and add types
