Improve the code

Made some changes to improve stuffs