Refactoring

and small bug fixies