Fix bugs and do improvements