Cleaned up the code after f*cking up the commit.
Needs further testing