improved code
wannabe fixes