Bug fixing and cleanup