Updated more code (sorry for lack of explanation haha..)