Javafx: How to implement 3-state checkboxes inside a TreeTableView

I am trying to implement a TreeTableView in javafx where the first column holds string values and the third one is to be rendered as 3-state checkboxes. With the following MCVE I am able to get a treetable but none of the selections in the checkboxes persist on parent collapse/expand or on resize of the table.
MCVE
Class A is parent.
Class B extends A and is child.
Class C represents the 2nd column, (rendered as checkboxes)
Class MVCECheckBox builds the treetable and displays it.

A.java

package mcve.checkbox;

import java.util.List;
import javafx.beans.property.SimpleStringProperty;
import javafx.beans.property.StringProperty;

/**
 *
 * @author returncode13
 */
public class A {
    StringProperty name=new SimpleStringProperty();
    C check=new C();
    List<A> children;

    public StringProperty getName() {
        return name;
    }

     public void setName(String name) {
        this.name.set(name);
    }

    public C getCheck() {
        return check;
    }

    public void setCheck(C check) {
        this.check = check;
    }

    public StringProperty nameProperty(){
        return name;
    }

    public List<A> getChildren() {
        return children;
    }

    public void setChildren(List<A> children) {
        this.children = children;
    }


    }

B.java

package mcve.checkbox;

import javafx.beans.property.SimpleStringProperty;
import javafx.beans.property.StringProperty;

/**
 *
 * @author returncode13
 */
public class B extends A{
    StringProperty name=new SimpleStringProperty();
    C Check=new C();

    @Override
    public StringProperty getName() {
        return name;
    }

    @Override
    public void setName(String name) {
        this.name.set(name);
    }

    @Override
    public C getCheck() {
        return Check;
    }

    @Override
    public void setCheck(C Check) {
        this.Check = Check;
    }

    @Override
     public StringProperty nameProperty(){
        return name;
    }

}

C.java

package mcve.checkbox;

import javafx.beans.property.BooleanProperty;
import javafx.beans.property.SimpleBooleanProperty;

/**
 *
 * @author returncode13
 */
public class C {
    BooleanProperty checkUncheck=new SimpleBooleanProperty();
    BooleanProperty indeterminate=new SimpleBooleanProperty();

    public BooleanProperty getCheckUncheck() {
        return checkUncheck;
    }

    public void setCheckUncheck(BooleanProperty checkUncheck) {
        this.checkUncheck = checkUncheck;
    }

    public BooleanProperty getIndeterminate() {
        return indeterminate;
    }

    public void setIndeterminate(BooleanProperty indeterminate) {
        this.indeterminate = indeterminate;
    }

    public BooleanProperty checkUncheckProperty(){
        return checkUncheck;
    }

    public BooleanProperty indeterminateProperty(){
        return indeterminate;
    }
} 

MCVECheckBox.java

package mcve.checkbox;
import javafx.beans.property.SimpleBooleanProperty;
import javafx.beans.value.ObservableValue;
import javafx.event.ActionEvent;
import javafx.event.EventHandler;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.control.CheckBox;
import javafx.scene.control.TreeItem;
import javafx.scene.control.TreeTableCell;
import javafx.scene.control.TreeTableColumn;
import javafx.scene.control.TreeTableView;
import javafx.scene.control.cell.TreeItemPropertyValueFactory;
import javafx.scene.layout.StackPane;
import javafx.stage.Stage;
import javafx.util.Callback;

/**
 *
 * @author returncode13
 */
public class MCVECheckBox extends Application {
    A selectedItem;
    private TreeTableView<A> treetable=new TreeTableView<>();
    @Override
    public void start(Stage primaryStage) {
       //setting up parents (A) and children (B)
        A a1=new A();                       
        a1.setName("A1");

        List<A> A1Children=new ArrayList();
        B b11=new B();
        b11.setName("B11");
        B b12=new B();
        b12.setName("B12");
        A1Children.add(b11);
        A1Children.add(b12);
        a1.setChildren(A1Children);


        A a2=new A();
        a2.setName("A2");

        List<A> A2Children=new ArrayList();
        B b21=new B();
        b21.setName("B21");
        B b22=new B();
        b22.setName("B22");
        A2Children.add(b21);
        A2Children.add(b22);
        a2.setChildren(A2Children);



        //tree columns . first one holds strings
        TreeTableColumn<A,String> name=new TreeTableColumn<>("Name");
        name.setCellValueFactory(new TreeItemPropertyValueFactory<>("name"));


        //2nd tree columns. rendered as checkboxes. boolean values
        TreeTableColumn<A,Boolean> checks=new TreeTableColumn<>("Checks");
        checks.setCellValueFactory(new Callback<TreeTableColumn.CellDataFeatures<A, Boolean>, ObservableValue<Boolean>>() {
            @Override
            public ObservableValue<Boolean> call(TreeTableColumn.CellDataFeatures<A, Boolean> param) {
                A a=param.getValue().getValue();
                SimpleBooleanProperty checkUncheck=new SimpleBooleanProperty();
                SimpleBooleanProperty indeterminate=new SimpleBooleanProperty();

                checkUncheck=(SimpleBooleanProperty) a.getCheck().getCheckUncheck();
                indeterminate=(SimpleBooleanProperty) a.getCheck().getIndeterminate();

                //to do: set parents status based on children status.


                if(indeterminate.get()){
                    return indeterminate;
                }else{
                    return checkUncheck;
                }

            }
        });


        checks.setCellFactory(new Callback<TreeTableColumn<A, Boolean>, TreeTableCell<A, Boolean>>() {
            @Override
            public TreeTableCell<A, Boolean> call(TreeTableColumn<A, Boolean> param) {
                return new CheckBoxCell();
            }
        });


        //building the tree;
        TreeItem<A> a1item=new TreeItem<>(a1);
        TreeItem<A> b11item=new TreeItem<>(b11);
        TreeItem<A> b12item=new TreeItem<>(b12);
        a1item.getChildren().add(b11item);
        a1item.getChildren().add(b12item);

        TreeItem<A> a2item=new TreeItem<>(a2);
        TreeItem<A> b21item=new TreeItem<>(b21);
        TreeItem<A> b22item=new TreeItem<>(b22);
        a2item.getChildren().add(b21item);
        a2item.getChildren().add(b22item);




        TreeItem<A> root=new TreeItem<>();
        root.getChildren().add(a1item);
        root.getChildren().add(a2item);

        treetable.getColumns().addAll(name,checks);

        treetable.setRoot(root);
        treetable.setShowRoot(false);
        treetable.setEditable(true);


      //  StackPane rootSp = new StackPane();
        Scene scene = new Scene(treetable, 300, 250);

        primaryStage.setTitle("Hello World!");
        primaryStage.setScene(scene);
        primaryStage.show();
    }

    /**
     * @param args the command line arguments
     */
    public static void main(String[] args) {
        launch(args);
    }


    //to render checkboxes in treetable


    private  class CheckBoxCell extends TreeTableCell<A, Boolean> {


        CheckBox checkbox;

        public CheckBoxCell() {
            checkbox=new CheckBox();
            checkbox.setAllowIndeterminate(true);

            checkbox.selectedProperty().addListener((obs,wasSelected,isNowSelected) -> {
            if(isNowSelected){
               selectedItem=getTreeTableRow().getItem();
            }
            });

        }


        @Override
        public void updateItem(Boolean b,boolean empty){
            super.updateItem(b, empty);

            if(empty){
                setGraphic(null);
            }else{
                checkbox.setSelected(b);
                setGraphic(checkbox);
            }
        }
    }

}

I have earlier tried using the CheckTreeTableCell to set the cell factory on the second column, but soon found out that the CheckTreeTableCell doesn't support 3-state (check,uncheck,indeterminate) checkboxes. After which I tried implementing the above code. Although I am able to bring in 3-state checkboxes, I am unable to let their state persist. Each time a parent is collapsed/expanded the checks made on its children become unselected.
Thanks for any help on determining a fix.

I am now able to implement the 3-state checkbox with the following modifications to the posted MCVE which is now a complete working example .

A.java (parent class)

package com.mycompany.yetanothercheckbox;

import java.util.Iterator;
import java.util.List;
import javafx.beans.property.SimpleStringProperty;
import javafx.beans.property.StringProperty;

/**
 *
 * @author returncode13
 */
public class A {
    private StringProperty name=new SimpleStringProperty();
    C check=new C();
    List<A> children;
    final boolean isLeaf=false;
    final boolean isParent=true;
    public boolean updateParent=false;
    public boolean updateChildren=false;
    public boolean isLeaf() {
        return isLeaf;
    }

    public boolean isParent() {
        return isParent;
    }
    public StringProperty getName() {
        return name;
    }

     public void setName(String name) {
        this.name.set(name);
    }

    public C getCheck() {
        return check;
    }

    public void setCheck(C check) {
        this.check = check;
        for (Iterator<A> iterator = children.iterator(); iterator.hasNext();) {
            A next = iterator.next();
            next.setCheck(check);

        }
    }

    public StringProperty nameProperty(){
        return name;
    }

    public List<A> getChildren() {
        return children;
    }

    public void setChildren(List<A> children) {
        this.children = children;
    }

    public A getParent() {
        return this;
    }
}

B.java (child class)

package com.mycompany.yetanothercheckbox;

import java.util.List;
import javafx.beans.property.SimpleStringProperty;
import javafx.beans.property.StringProperty;

/**
 *
 * @author returncode13
 */
public class B extends A{
    private StringProperty name=new SimpleStringProperty();
    C Check=new C();
    final boolean isLeaf=true;
    final boolean isParent=false;
     public boolean updateParent=false;
    public boolean updateChildren=false;
    A parent;

    public A getParent() {
        return parent;
    }

    public void setParent(A parent) {
        this.parent = parent;
    }


    @Override
    public boolean isLeaf() {
        return isLeaf;
    }

    @Override
    public boolean isParent() {
        return isParent;
    }



    @Override
    public StringProperty getName() {
        return name;
    }

    @Override
    public void setName(String name) {
        this.name.set(name);
    }

    @Override
    public C getCheck() {
        return Check;
    }

    @Override
    public void setCheck(C Check) {
        this.Check = Check;
    }

    @Override
     public StringProperty nameProperty(){
        return name;
    }

    @Override
       public List<A> getChildren() {
        return parent.getChildren();
    }

}

C.java (Hold Check states)

package com.mycompany.yetanothercheckbox;

import javafx.beans.property.BooleanProperty;
import javafx.beans.property.SimpleBooleanProperty;

/**
 *
 * @author returncode13
 */
public class C {
    BooleanProperty checkUncheck=new SimpleBooleanProperty();
    BooleanProperty indeterminate=new SimpleBooleanProperty();

    public BooleanProperty getCheckUncheck() {
        return checkUncheck;
    }

    public void setCheckUncheck(BooleanProperty checkUncheck) {
        this.checkUncheck = checkUncheck;
    }

    public BooleanProperty getIndeterminate() {
        return indeterminate;
    }

    public void setIndeterminate(BooleanProperty indeterminate) {
        this.indeterminate = indeterminate;
    }

    public BooleanProperty checkUncheckProperty(){
        return checkUncheck;
    }

    public BooleanProperty indeterminateProperty(){
        return indeterminate;
    }
}

ThreeStateCheckBoxTreeTableCell.java (the 3state checkbox for tree table)

package com.mycompany.yetanothercheckbox;

import java.util.List;
import javafx.beans.value.ChangeListener;
import javafx.beans.value.ObservableValue;
import javafx.event.EventHandler;
import javafx.scene.control.CheckBox;
import javafx.scene.control.TreeTableCell;
import javafx.scene.control.TreeTableColumn;
import javafx.scene.input.MouseEvent;

/**
 *
 * @author returncode13
 */
   //to render checkboxes in treetable


    public  class ThreeStateCheckBoxTreeTableCell extends TreeTableCell<A, Boolean> {


        A selectedItem;
        CheckBox checkbox;
        TreeTableColumn<A,Boolean> param;
        /*private static boolean updateParent=false;
        private static boolean updateChildren=false;*/


        public ThreeStateCheckBoxTreeTableCell(TreeTableColumn<A,Boolean> param) {
            checkbox=new CheckBox();
            this.param=param;
            checkbox.setAllowIndeterminate(true);

            checkbox.selectedProperty().addListener((obs,wasSelected,isNowSelected) -> {
               int sel=getTreeTableRow().getIndex();
               selectedItem=this.param.getTreeTableView().getSelectionModel().getModelItem(sel).getValue();
               selectedItem.getCheck().getCheckUncheck().set(isNowSelected);
               selectedItem.getCheck().getIndeterminate().set(false); 

              //ThreeStateCheckBoxTreeTableCell.this.param.getTreeTableView().refresh();


            });


            checkbox.indeterminateProperty().addListener((obx,ol,newV)->{

               int sel=getTreeTableRow().getIndex();
               selectedItem=this.param.getTreeTableView().getSelectionModel().getModelItem(sel).getValue();
                selectedItem.getCheck().getIndeterminate().set(newV);

           //ThreeStateCheckBoxTreeTableCell.this.param.getTreeTableView().refresh();
            });



            checkbox.setOnMouseClicked(new EventHandler<MouseEvent>(){
                @Override
                public void handle(MouseEvent event) {
                    int sel=getTreeTableRow().getIndex();
                    selectedItem=ThreeStateCheckBoxTreeTableCell.this.param.getTreeTableView().getSelectionModel().getModelItem(sel).getValue();
                    if(selectedItem.isParent()){

                        selectedItem.updateChildren=true;
                        for(A child:selectedItem.getChildren()){
                            child.updateParent=false;
                        }


                         updateDownwards();
                    }
                    if(selectedItem.isLeaf()){

                        selectedItem.getParent().updateChildren=false;
                        selectedItem.updateParent=true;

                        updateUpWards();
                    }

                    ThreeStateCheckBoxTreeTableCell.this.param.getTreeTableView().refresh();
                }

            });
        }


        @Override
        public void updateItem(Boolean b,boolean empty){
            super.updateItem(b, empty);

            if(empty){
                setGraphic(null);
            }else{


                if(b==null){
                    checkbox.setIndeterminate(true);

                }
                else{
                    checkbox.setIndeterminate(false);
                    checkbox.setSelected(b);

                }
               setGraphic(checkbox);

            }

            ThreeStateCheckBoxTreeTableCell.this.param.getTreeTableView().refresh();
        }

        private void updateUpWards(){

            if(selectedItem.updateParent){


            List<A> children=selectedItem.getChildren();
            int indeterminateCount=0;
            int selectedCount=0;

                A parent=selectedItem.getParent();
                for(A child:children){
                    indeterminateCount+=child.getCheck().getIndeterminate().get()?1:0;
                    selectedCount+=child.getCheck().getCheckUncheck().get()?1:0;

                }

                if(indeterminateCount>0) {

                     parent.getCheck().getIndeterminate().set(true);
                }
                else if(indeterminateCount==0 && selectedCount==children.size()){
                    parent.getCheck().getIndeterminate().set(false);
                    parent.getCheck().getCheckUncheck().set(true);
                }else{

                    parent.getCheck().getIndeterminate().set(false);
                    parent.getCheck().getCheckUncheck().set(false);
                }

            }
            ThreeStateCheckBoxTreeTableCell.this.param.getTreeTableView().refresh();


        }


        private void updateDownwards(){
             List<A> children=selectedItem.getChildren();

             if(selectedItem.isParent() && selectedItem.updateChildren ){

                 for(A child:children){
                   child.getCheck().getCheckUncheck().set(selectedItem.getCheck().getCheckUncheck().get());

                  child.getCheck().getIndeterminate().set(selectedItem.getCheck().getIndeterminate().get());
                }


             }
             ThreeStateCheckBoxTreeTableCell.this.param.getTreeTableView().refresh();

        }
    }

MainApp.java (Application as a POC)

package com.mycompany.yetanothercheckbox;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import javafx.application.Application;
import static javafx.application.Application.launch;
import javafx.beans.property.SimpleBooleanProperty;
import javafx.beans.value.ChangeListener;
import javafx.beans.value.ObservableValue;
import javafx.fxml.FXMLLoader;
import javafx.scene.Parent;
import javafx.scene.Scene;
import javafx.scene.control.CheckBox;
import javafx.scene.control.TreeItem;
import javafx.scene.control.TreeTableCell;
import javafx.scene.control.TreeTableColumn;
import javafx.scene.control.TreeTableView;
import javafx.scene.control.cell.TreeItemPropertyValueFactory;
import javafx.stage.Stage;
import javafx.util.Callback;


public class MainApp extends Application {

    private TreeTableView<A> treetable=new TreeTableView<>();
    @Override
    public void start(Stage primaryStage) {
       //setting up parents (A) and children (B)
        A a1=new A();                       
        a1.setName("A1");

        List<A> A1Children=new ArrayList();
        B b11=new B();
        b11.setName("B11");
        B b12=new B();
        b12.setName("B12");
        b11.setParent(a1);
        b12.setParent(a1);
        A1Children.add(b11);
        A1Children.add(b12);
        a1.setChildren(A1Children);


        A a2=new A();
        a2.setName("A2");

        List<A> A2Children=new ArrayList();
        B b21=new B();
        b21.setName("B21");
        B b22=new B();
        b22.setName("B22");
        b21.setParent(a2);
        b22.setParent(a2);
        A2Children.add(b21);
        A2Children.add(b22);
        a2.setChildren(A2Children);



        //tree columns . first one holds strings
        TreeTableColumn<A,String> name=new TreeTableColumn<>("Name");
        name.setCellValueFactory(new TreeItemPropertyValueFactory<>("name"));


        //2nd tree columns. rendered as checkboxes. boolean values
        TreeTableColumn<A,Boolean> checks=new TreeTableColumn<>("Checks");
        checks.setCellValueFactory(new Callback<TreeTableColumn.CellDataFeatures<A, Boolean>, ObservableValue<Boolean>>() {
            @Override
            public ObservableValue<Boolean> call(TreeTableColumn.CellDataFeatures<A, Boolean> param) {
                A a=param.getValue().getValue();
                SimpleBooleanProperty checkUncheck=new SimpleBooleanProperty();
                SimpleBooleanProperty indeterminate=new SimpleBooleanProperty();

                checkUncheck.bindBidirectional(a.getCheck().getCheckUncheck());
                indeterminate.bindBidirectional(a.getCheck().getIndeterminate());


                checkUncheck.addListener(new ChangeListener<Boolean>(){
                    @Override
                    public void changed(ObservableValue<? extends Boolean> observable, Boolean oldValue, Boolean newValue) {

                        a.getCheck().indeterminateProperty().set(false);
                        a.getCheck().checkUncheckProperty().set(newValue);
                    }

                });
                indeterminate.addListener(new ChangeListener<Boolean>(){
                    @Override
                    public void changed(ObservableValue<? extends Boolean> observable, Boolean oldValue, Boolean newValue) {

                            a.getCheck().indeterminateProperty().set(newValue);

                    }
                });
                if(indeterminate.get()){

                return null;
                }else{

                    return checkUncheck;
                }

            }
        });


        checks.setCellFactory(new Callback<TreeTableColumn<A, Boolean>, TreeTableCell<A, Boolean>>() {
            @Override
            public TreeTableCell<A, Boolean> call(TreeTableColumn<A, Boolean> param) {
                return new ThreeStateCheckBoxTreeTableCell(param);
            }
        });


        //building the tree;
        TreeItem<A> a1item=new TreeItem<>(a1);
        TreeItem<A> b11item=new TreeItem<>(b11);
        TreeItem<A> b12item=new TreeItem<>(b12);
        a1item.getChildren().add(b11item);
        a1item.getChildren().add(b12item);

        TreeItem<A> a2item=new TreeItem<>(a2);
        TreeItem<A> b21item=new TreeItem<>(b21);
        TreeItem<A> b22item=new TreeItem<>(b22);
        a2item.getChildren().add(b21item);
        a2item.getChildren().add(b22item);




        TreeItem<A> root=new TreeItem<>();
        root.getChildren().add(a1item);
        root.getChildren().add(a2item);

        treetable.getColumns().addAll(name,checks);

        treetable.setRoot(root);
        treetable.setShowRoot(false);
        treetable.setEditable(true);


      //  StackPane rootSp = new StackPane();
        Scene scene = new Scene(treetable, 300, 250);

        primaryStage.setTitle("CheckBoxTreeTable Example");
        primaryStage.setScene(scene);
        primaryStage.show();
    }

    /**
     * @param args the command line arguments
     */
    public static void main(String[] args) {
        launch(args);
    }




}