Added some edge cases and correted the cleaning func