clean up/bug fix