import java.util.ArrayList;
import javafx.application.Application;
import static javafx.application.Application.launch;
import javafx.beans.property.SimpleStringProperty;
import javafx.beans.value.ObservableValue;
import javafx.collections.ObservableList;
import javafx.event.ActionEvent;
import javafx.event.EventHandler;
import javafx.scene.Node;
import javafx.scene.Scene;
import javafx.scene.control.TreeItem;
import javafx.scene.control.TreeTableColumn;
import javafx.scene.control.TreeTableView;
import javafx.scene.layout.HBox;
import javafx.scene.layout.Pane;
import javafx.scene.layout.VBox;
import javafx.util.Callback;
import javafx.stage.Stage;
import javafx.scene.control.Button;

public class TreeTableTest extends Application {

    public static final int DATA_FIELDS_NUM = 2;
    public static final int DATA_ITEMS_NUM = 5;
    Pane content;
    TreeTableView<Data> treeTableView;
    ObservableList<TreeItem<Data>> data;

    public static void main(String args[]) {
        launch(TreeTableTest.class, args);
    }

    @Override
    public void start(Stage primaryStage) {
        HBox cont = new HBox(2);
        Scene scene = new Scene(cont);
        cont.getChildren().add(getTestContainer());
        cont.getChildren().add(getControlsContainer());
        primaryStage.setScene(scene);
        primaryStage.setTitle("TreeTable Test " + System.getProperty("javafx.runtime.version"));
        primaryStage.show();
    }

    protected Node getTestContainer() {
        content = new Pane();
        content.setPrefSize(500, 500);
        treeTableView = new TreeTableView<Data>();
        treeTableView.setTranslateX(50);
        treeTableView.setTranslateY(50);
        treeTableView.setPrefSize(400, 400);

        reset(DATA_ITEMS_NUM, DATA_FIELDS_NUM);
        content.getChildren().add(treeTableView);

        return content;
    }

    protected void reset(int rows, int columns) {
        TreeItem root = new TreeItem(new Data("root", columns));
        root.setExpanded(true);
        treeTableView.setRoot(root);
        data = treeTableView.getRoot().getChildren();
        data.clear();
        for (int i = 0; i < rows; i++) {
            data.add(new TreeItem(new Data(String.format("%02d", i), columns)));
        }
        treeTableView.getColumns().clear();
        for (int i = 0; i < columns; i++) {
            treeTableView.getColumns().add(createColumn(i));
        }
    }

    protected TreeTableColumn<Data, String> createColumn(final int index) {
        final TreeTableColumn<Data, String> field = new TreeTableColumn<Data, String>("field " + index);

        field.setCellValueFactory(new Callback<TreeTableColumn.CellDataFeatures<Data, String>, ObservableValue<String>>() {
            public ObservableValue<String> call(TreeTableColumn.CellDataFeatures<Data, String> p) {
                return ((Data) ((TreeItem) p.getValue()).getValue()).data.get(index);
            }
        });
        return field;
    }

    protected Node getControlsContainer() {
        VBox res = new VBox(10);
        res.setTranslateY(50);

        Button reset = new Button("reset");
        reset.setId("reset");
        reset.setOnAction(new EventHandler<ActionEvent>() {
            public void handle(ActionEvent t) {
                reset(DATA_ITEMS_NUM, DATA_FIELDS_NUM);
            }
        });
        res.getChildren().add(reset);

        Button btnRemSelected = new Button("remove selected");
        btnRemSelected.setOnAction(new EventHandler<ActionEvent>() {
            public void handle(ActionEvent t) {
                ObservableList<TreeItem<Data>> selectedItems
                        = treeTableView.getSelectionModel().getSelectedItems();
                ArrayList<TreeItem<Data>> list = new ArrayList<TreeItem<Data>>();
                for (TreeItem<Data> item : selectedItems) {
                    list.add(item);
                }
                if (list.contains(treeTableView.getRoot())) {
                    System.out.println("root is contained in selected items and going to be deleted!!");
                    treeTableView.setRoot(null);
                } else {
                    for (TreeItem<Data> item : list) {
                        remove(treeTableView.getRoot(), item);
                    }
                }
            }
        });
        res.getChildren().add(btnRemSelected);

        return res;
    }

    public static class Data {

        protected String name;
        protected ArrayList<SimpleStringProperty> data;

        public Data(String name, int fieldsNum) {
            this.name = name;
            data = new ArrayList<SimpleStringProperty>(fieldsNum);
            for (int i = 0; i < fieldsNum; i++) {
                data.add(new SimpleStringProperty("item " + name + " field " + i));
            }
        }
    }

    protected void remove(TreeItem root, TreeItem item) {
        if (!root.getChildren().remove(item)) {
            for (TreeItem<Data> child : (ObservableList<TreeItem<Data>>) root.getChildren()) {
                remove(child, item);
            }
        }
    }
}