/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.examples.cookbook;

import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.Logger;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.Validation;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.Element;
import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
import org.apache.beam.sdk.transforms.DoFn.SideInput;
import org.apache.beam.sdk.transforms.Mean;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;

/**
 * This is an example that demonstrates several approaches to filtering, and use of the Mean
 * transform. It shows how to dynamically set parameters by defining and using new pipeline options,
 * and how to use a value derived by the pipeline.
 *
 * <p>Concepts: The Mean transform; Options configuration; using pipeline-derived data as a side
 * input; approaches to filtering, selection, and projection.
 *
 * <p>The example reads public samples of weather data from BigQuery. It performs a projection on
 * the data, finds the global mean of the temperature readings, filters on readings for a single
 * given month, and then outputs only data (for that month) that has a mean temp smaller than the
 * derived global mean.
 *
 * <p>Note: Before running this example, you must create a BigQuery dataset to contain your output
 * table.
 *
 * <p>To execute this pipeline locally, specify the BigQuery table for the output:
 *
 * <pre>{@code
 * --output=YOUR_PROJECT_ID:DATASET_ID.TABLE_ID
 * [--monthFilter=<month_number>]
 * }</pre>
 *
 * where optional parameter {@code --monthFilter} is set to a number 1-12.
 *
 * <p>To change the runner, specify:
 *
 * <pre>{@code
 * --runner=YOUR_SELECTED_RUNNER
 * }</pre>
 *
 * See examples/java/README.md for instructions about how to configure different runners.
 *
 * <p>The BigQuery input table defaults to {@code apache-beam-testing.samples.weather_stations} and
 * can be overridden with {@code --input}.
 */
public class FilterExamples {
  // Default to using a 1000 row subset of the public weather station table publicdata:samples.gsod.
  private static final String WEATHER_SAMPLES_TABLE =
      "apache-beam-testing.samples.weather_stations";
  static final Logger LOG = Logger.getLogger(FilterExamples.class.getName());
  static final int MONTH_TO_FILTER = 7;

  /**
   * Examines each row in the input table. Outputs only the subset of the cells this example is
   * interested in-- the mean_temp and year, month, and day-- as a bigquery table row.
   */
  static class ProjectionFn extends DoFn<TableRow, TableRow> {
    @ProcessElement
    public void processElement(@Element TableRow row, OutputReceiver<TableRow> receiver) {
      // Grab year, month, day, mean_temp from the row
      Integer year = Integer.parseInt((String) row.get("year"));
      Integer month = Integer.parseInt((String) row.get("month"));
      Integer day = Integer.parseInt((String) row.get("day"));
      Double meanTemp = Double.parseDouble(row.get("mean_temp").toString());
      // Prepares the data for writing to BigQuery by building a TableRow object
      TableRow outRow =
          new TableRow()
              .set("year", year)
              .set("month", month)
              .set("day", day)
              .set("mean_temp", meanTemp);
      receiver.output(outRow);
    }
  }

  /**
   * Implements 'filter' functionality.
   *
   * <p>Examines each row in the input table. Outputs only rows from the month monthFilter, which is
   * passed in as a parameter during construction of this DoFn.
   */
  static class FilterSingleMonthDataFn extends DoFn<TableRow, TableRow> {
    Integer monthFilter;

    public FilterSingleMonthDataFn(Integer monthFilter) {
      this.monthFilter = monthFilter;
    }

    @ProcessElement
    public void processElement(@Element TableRow row, OutputReceiver<TableRow> receiver) {
      Integer month;
      month = (Integer) row.get("month");
      if (month.equals(this.monthFilter)) {
        receiver.output(row);
      }
    }
  }

  /**
   * Examines each row (weather reading) in the input table. Output the temperature reading for that
   * row ('mean_temp').
   */
  static class ExtractTempFn extends DoFn<TableRow, Double> {
    @ProcessElement
    public void processElement(@Element TableRow row, OutputReceiver<Double> receiver) {
      Double meanTemp = Double.parseDouble(row.get("mean_temp").toString());
      receiver.output(meanTemp);
    }
  }

  /**
   * Finds the global mean of the mean_temp for each day/record, and outputs only data that has a
   * mean temp larger than this global mean.
   */
  static class BelowGlobalMean extends PTransform<PCollection<TableRow>, PCollection<TableRow>> {
    Integer monthFilter;

    public BelowGlobalMean(Integer monthFilter) {
      this.monthFilter = monthFilter;
    }

    @Override
    public PCollection<TableRow> expand(PCollection<TableRow> rows) {

      // Extract the mean_temp from each row.
      PCollection<Double> meanTemps = rows.apply(ParDo.of(new ExtractTempFn()));

      // Find the global mean, of all the mean_temp readings in the weather data,
      // and prepare this singleton PCollectionView for use as a side input.
      final PCollectionView<Double> globalMeanTemp =
          meanTemps.apply(Mean.globally()).apply(View.asSingleton());

      // Rows filtered to remove all but a single month
      PCollection<TableRow> monthFilteredRows =
          rows.apply(ParDo.of(new FilterSingleMonthDataFn(monthFilter)));

      // Then, use the global mean as a side input, to further filter the weather data.
      // By using a side input to pass in the filtering criteria, we can use a value
      // that is computed earlier in pipeline execution.
      // We'll only output readings with temperatures below this mean.
      PCollection<TableRow> filteredRows =
          monthFilteredRows.apply(
              "ParseAndFilter",
              ParDo.of(
                      new DoFn<TableRow, TableRow>() {
                        @ProcessElement
                        public void processElement(
                            @SideInput("globalMeanTemp") Double gTemp,
                            @Element TableRow element,
                            OutputReceiver<TableRow> receiver) {
                          Double meanTemp = Double.parseDouble(element.get("mean_temp").toString());
                          if (meanTemp < gTemp) {
                            receiver.output(element);
                          }
                        }
                      })
                  .withSideInput("globalMeanTemp", globalMeanTemp));

      return filteredRows;
    }
  }

  /**
   * Options supported by {@link FilterExamples}.
   *
   * <p>Inherits standard configuration options.
   */
  public interface Options extends PipelineOptions {
    @Description("Table to read from, specified as " + "<project_id>:<dataset_id>.<table_id>")
    @Default.String(WEATHER_SAMPLES_TABLE)
    String getInput();

    void setInput(String value);

    @Description(
        "Table to write to, specified as "
            + "<project_id>:<dataset_id>.<table_id>. "
            + "The dataset_id must already exist")
    @Validation.Required
    String getOutput();

    void setOutput(String value);

    @Description("Numeric value of month to filter on")
    @Default.Integer(MONTH_TO_FILTER)
    Integer getMonthFilter();

    void setMonthFilter(Integer value);
  }

  /** Helper method to build the table schema for the output table. */
  private static TableSchema buildWeatherSchemaProjection() {
    List<TableFieldSchema> fields = new ArrayList<>();
    fields.add(new TableFieldSchema().setName("year").setType("INTEGER"));
    fields.add(new TableFieldSchema().setName("month").setType("INTEGER"));
    fields.add(new TableFieldSchema().setName("day").setType("INTEGER"));
    fields.add(new TableFieldSchema().setName("mean_temp").setType("FLOAT"));
    return new TableSchema().setFields(fields);
  }

  public static void main(String[] args) throws Exception {

    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
    runFilterExamples(options);
  }

  static void runFilterExamples(Options options) {
    Pipeline p = Pipeline.create(options);

    TableSchema schema = buildWeatherSchemaProjection();

    p.apply(BigQueryIO.readTableRows().from(options.getInput()))
        .apply(ParDo.of(new ProjectionFn()))
        .apply(new BelowGlobalMean(options.getMonthFilter()))
        .apply(
            BigQueryIO.writeTableRows()
                .to(options.getOutput())
                .withSchema(schema)
                .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
                .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE));

    p.run().waitUntilFinish();
  }
}
