bug fixing and clean up