Fix bugs, make code nicer