Bug fixing/cleanup