fixed a bug
Refactored the code