Code refactored and bugs fixed (?)