Bug fix
Cleanup
Added checks