Bug fixes and more improvments