fix some bug and make code better