Fixed one error which made it break the test suite. Now it doesn't break, but still needs more test cases to see if the functionality works on all inputs.