Fix small bugs and clean up code