Bug fixes and improvments