Tables.java

//Dstl (c) Crown Copyright 2017
package uk.gov.dstl.baleen.annotators.structural;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.jcas.JCas;
import org.apache.uima.resource.ResourceInitializationException;

import uk.gov.dstl.baleen.types.structure.Structure;
import uk.gov.dstl.baleen.types.structure.Table;
import uk.gov.dstl.baleen.types.structure.TableCell;
import uk.gov.dstl.baleen.uima.utils.StructureHierarchy;
import uk.gov.dstl.baleen.uima.utils.StructureUtil;
import uk.gov.dstl.baleen.uima.utils.select.ItemHierarchy;
import uk.gov.dstl.baleen.uima.utils.select.Node;
import uk.gov.dstl.baleen.uima.utils.select.Nodes;

/**
 * Helper class for working with {@link Table}s.
 *
 * <p>
 * This is an example of how structure selectors can be used
 */
public class Tables {

	private List<Pattern> columnPatterns = new ArrayList<>();
	private Nodes<Structure> nodes;

	/**
	 * Constructor for tables helper
	 *
	 * @param jCas
	 *            the jcas
	 * @throws AnalysisEngineProcessException
	 *             if class creation error
	 */
	public Tables(JCas jCas) throws AnalysisEngineProcessException {
		try {
			ItemHierarchy<Structure> hierarchy = StructureHierarchy.build(jCas, StructureUtil.getStructureClasses());
			nodes = hierarchy.getRoot().select("Table");
		} catch (ResourceInitializationException e) {
			throw new AnalysisEngineProcessException("Can not create structure helper", null, e);
		}
	}

	/**
	 * Filter to tables with a column matching the given name.
	 * <p>
	 * This uses the table header to find column names.
	 *
	 *
	 * @param p
	 *            the pattern to match
	 * @return this for builder pattern
	 */
	public Tables withColumn(String p) {
		return withColumn(Pattern.compile(p));
	}

	/**
	 * Filter to tables with a column matching the given name.
	 * <p>
	 * This uses the table header to find column names.
	 *
	 *
	 * @param p
	 *            the pattern to match
	 * @return this for builder pattern
	 */
	public Tables withColumn(Pattern p) {
		nodes = nodes.select(":has(TableHeader:matches(" + p.pattern() + "))");
		columnPatterns.add(p);
		return this;
	}

	/**
	 * Get the rows of the tables filtered to the specified columns and in given
	 * column order.
	 *
	 * @return a stream of the filtered rows (ie list of TableCells)
	 */
	public Stream<List<TableCell>> getFilteredRows() {
		return nodes.stream().flatMap(table -> {

			List<Nodes<Structure>> columns = columnPatterns.stream()
					.map(p -> nodes.select("TableHeader TableCell:matches(" + p.pattern() + ")"))
					.map(p -> p.get(0))
					.map(Node::getSiblingIndex)
					.map(i -> nodes.select("TableBody > TableRow > TableCell:nth-child(" + (i + 1) + ")"))
					.collect(Collectors.toList());

			if (columns.isEmpty()) {
				return Stream.empty();
			}
			// validate that the input lists are all the same size.
			int numItems = columns.get(0).size();
			for (int i = 1; i < columns.size(); i++) {
				if (columns.get(i).size() != numItems) {
					// non-uniform-length list at index i
					return Stream.empty();
				}
			}

			List<List<TableCell>> result = new ArrayList<>();

			for (int i = 0; i < numItems; i++) {

				// create a tuple of the i-th entries of each list
				List<TableCell> row = new ArrayList<>(columns.size());
				for (Nodes<Structure> column : columns) {
					row.add((TableCell) column.get(i).getItem());
				}
				result.add(row);
			}

			return result.stream();

		});

	}

	/**
	 * Get the cells of the tables filtered to the specified columns.
	 *
	 * @return a stream of the filtered table cells
	 */
	public Stream<TableCell> getFilteredCells() {
		return getFilteredRows().flatMap(List::stream);
	}

	/**
	 * Get the tables which pass the specified column filters.
	 *
	 * @return a stream of the filtered tables
	 */
	public Stream<Table> getTables() {
		return nodes.stream().map(t -> (Table) t.getItem());
	}

	/**
	 * Get the full rows of the tables which pass the specified column filters.
	 *
	 * @return a stream of the filtered table cells
	 */
	public Stream<List<TableCell>> getRows() {
		return nodes.select("TableBody > TableRow").stream()
				.map(tr -> tr.getChildren().stream().map(tc -> (TableCell) tc.getItem()).collect(Collectors.toList()));
	}

	/**
	 * Get all the cells of the tables which pass the specified column filters.
	 *
	 * @return a stream of the table cells
	 */
	public Stream<TableCell> getCells() {
		return getRows().flatMap(List::stream);
	}

}