Code cleanup + bugfix
Fix stupid error