Fix bugs and new improvements