/**
 * A Graph acts as an Estimator. It consists of a DAG of stages, each of which is either an
 * Estimator or Transformer.
 */
@PublicEvolving
public final class Graph implements Estimator<Graph, GraphModel> {
    public Graph(...) {...}

    @Override
    public GraphModel fit(Table... inputs) {...}

    @Override
    public TableSchema[] transformSchemas(TableSchema... schemas) {
        return schemas;
    }

    /** Skipped a few methods, including the implementations of some Estimator APIs. */
}


/** A GraphBuilder helps connect Stage instances into a Graph or GraphModel. */
@PublicEvolving
public final class GraphBuilder {
    private int maxOutputLength = 20;

    public GraphBuilder() {

    }

    /**
     * Specifies the upper bound (could be loose) of the number of output tables that can be returned by the
     * Transformer::getStateStreams and Transformer::transform methods, for any stage involved in this Graph.
     *
     * The default upper bound is 20.
     */
    public GraphBuilder setMaxOutputLength(int maxOutputLength) {
        this.maxOutputLength = maxOutputLength;
        return this;
    }

    /**
     * Creates a TableId associated with this GraphBuilder. It can be used to specify the passing of
     * tables between stages, as well as the input/output tables of the Graph/GraphModel generated by
     * this builder.
     */
    public TableId createTableId() {
        return new TableId();
    }

    /**
     * The Graph::fit and GraphModel::transform should invoke the fit/transform of the corresponding stage with the
     * corresponding inputs.
     *
     * Returns a list of TableIds, which represents outputs of the Transformer::transform invocation.
     */
    public TableId[] getOutputs(Stage<?> stage, TableId... inputs) {
        return new TableId[maxOutputLength];
    }

    /**
     * The GraphModel::setStateStreams should invoke the setStateStreams of the corresponding stage with the
     * corresponding inputs.
     */
    void setStateStreams(Stage<?> stage, TableId... inputs) {}


    /**
     * The GraphModel::getStateStreams should invoke the getStateStreams of the corresponding stage.
     *
     * Returns a list of TableIds, which represents outputs of the getStateStreams invocation.
     */
    TableId[] getStateStreams(Stage<?> stage) {
        return new TableId[maxOutputLength];
    }

    /**
     * Returns a Graph instance which the following API specification:
     * - Graph::fit should take inputs and returns a GraphModel with the following specification.
     * - GraphModel::transform should take inputs and returns outputs.
     * - GraphModel::setStateStreams should take inputStates.
     * - GraphModel::getStateStreams should return outputStates.
     *
     * The fit/transform/setStateStreams/getStateStreams should invoke the APIs of the internal stages in
     * the order specified by the DAG of stages.
     */
    Graph build(
            TableId[] inputs, TableId[] outputs, TableId[] inputStates, TableId[] outputStates) {
        return new Graph();
    }

    /**
     * Returns a GraphModel instance which the following API specification:
     * - GraphModel::transform should take inputs and returns outputs.
     * - GraphModel::setStateStreams should take inputStates.
     * - GraphModel::getStateStreams should return outputStates.
     *
     * The transform/setStateStreams/getStateStreams should invoke the APIs of the internal stages in
     * the order specified by the DAG of stages.
     *
     * This method throws exception if any stage of this graph is an Estimator.
     */
    GraphModel buildModel(
            TableId[] inputs, TableId[] outputs, TableId[] inputStates, TableId[] outputStates) {
        return new GraphModel();
    }

    static class TableId {}
}