Fix bugs and better code