Bug fixes and code cleaning