Fixed the bug and more clean up

see title