Refactored the code so it actually works... :)