Cleaning up functionality and fix some bugs for Java