Improvements

Made some changes, let me know