fixed some bugs and improve the code