Lots of changes after fix

All fixed and lots of refactoring is done