fixing bug and code clean up.