Fix bugs and clean code