Big fixes and code improvements