name: Prepare data for train
description: |
    Prepares the parameters for the training step.

    Converts the input_tables and the output of ForecastingPreprocessingOp
    to the input parameters of TimeSeriesDatasetCreateOp and
    AutoMLForecastingTrainingJobRunOp.

    Args:
      input_tables (JsonArray):
        Required. Serialized Json array that specifies input BigQuery tables and
        specs.
      preprocess_metadata (JsonObject):
        Required. The output of ForecastingPreprocessingOp that is a serialized
        dictionary with 2 fields: processed_bigquery_table_uri and
        column_metadata.
      model_feature_columns (JsonArray):
        Optional. Serialized list of column names that will be used as input
        feature in the training step. If None, all columns will be used in
        training.

    Returns:
      NamedTuple:
        time_series_identifier_column (String):
          Name of the column that identifies the time series.
        time_series_attribute_columns (JsonArray):
          Serialized column names that should be used as attribute columns.
        available_at_forecast_columns (JsonArray):
          Serialized column names of columns that are available at forecast.
        unavailable_at_forecast_columns (JsonArray):
          Serialized column names of columns that are unavailable at forecast.
        column_transformations (JsonArray):
          Serialized transformations to apply to the input columns.
        preprocess_bq_uri (String):
          The BigQuery table that saves the preprocessing result and will be
          used as training input.
        target_column (String):
          The name of the column values of which the Model is to predict.
        time_column (String):
          Name of the column that identifies time order in the time series.
        predefined_split_column (String):
          Name of the column that specifies an ML use of the row.
        weight_column (String):
          Name of the column that should be used as the weight column.
        data_granularity_unit (String):
          The data granularity unit.
        data_granularity_count (Integer):
          The number of data granularity units between data points in the
          training data.
inputs:
- name: input_tables
  type: JsonArray
  description: |-
    Required. Serialized Json array that specifies input BigQuery tables and
    specs.
- name: preprocess_metadata
  type: JsonObject
  description: |-
    Required. The output of ForecastingPreprocessingOp that is a serialized
    dictionary with 2 fields: processed_bigquery_table_uri and
    column_metadata.
- name: model_feature_columns
  type: JsonArray
  description: |-
    Optional. Serialized list of column names that will be used as input
    feature in the training step. If None, all columns will be used in
    training.
  optional: true
outputs:
- {name: time_series_identifier_column, type: String}
- {name: time_series_attribute_columns, type: JsonArray}
- {name: available_at_forecast_columns, type: JsonArray}
- {name: unavailable_at_forecast_columns, type: JsonArray}
- {name: column_transformations, type: JsonArray}
- {name: preprocess_bq_uri, type: String}
- {name: target_column, type: String}
- {name: time_column, type: String}
- {name: predefined_split_column, type: String}
- {name: weight_column, type: String}
- {name: data_granularity_unit, type: String}
- {name: data_granularity_count, type: Integer}
implementation:
  container:
    image: python:3.8
    command:
    - sh
    - -ec
    - |
      program_path=$(mktemp)
      printf "%s" "$0" > "$program_path"
      python3 -u "$program_path" "$@"
    - |
      def prepare_data_for_train(
          # pylint: disable=g-bare-generic
          input_tables,
          preprocess_metadata,
          model_feature_columns = None,
          # pylint: enable=g-bare-generic
      ):
        # pylint: disable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
        from typing import NamedTuple
        # pylint: enable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported

        column_metadata = preprocess_metadata['column_metadata']
        bigquery_table_uri = preprocess_metadata['processed_bigquery_table_uri']

        primary_table_specs = next(
            table for table in input_tables
            if table['table_type'] == 'FORECASTING_PRIMARY'
        )
        primary_metadata = primary_table_specs['forecasting_primary_table_metadata']

        feature_columns = None
        if model_feature_columns:
          feature_columns = set(model_feature_columns)

        time_series_identifier_column = ''
        time_series_attribute_columns = []
        available_at_forecast_columns = []
        unavailable_at_forecast_columns = []
        column_transformations = []
        predefined_split_column = (
            '' if 'predefined_splits_column' not in primary_metadata
            else primary_metadata['predefined_splits_column'])
        weight_column = (
            '' if 'weight_column' not in primary_metadata
            else primary_metadata['weight_column'])

        for name, details in column_metadata.items():
          if name == predefined_split_column or name == weight_column:
            continue
          # Determine temporal type of the column
          if details['tag'] == 'vertex_time_series_key':
            time_series_identifier_column = name
            continue
          elif details['tag'] == 'primary_table':
            if name in (
                primary_metadata['time_series_identifier_columns']
            ):
              time_series_attribute_columns.append(name)
            elif name == primary_metadata['target_column']:
              unavailable_at_forecast_columns.append(name)
            elif name in primary_metadata['unavailable_at_forecast_columns']:
              unavailable_at_forecast_columns.append(name)
            else:
              available_at_forecast_columns.append(name)
          elif details['tag'] == 'attribute_table':
            time_series_attribute_columns.append(name)
          elif details['tag'] == 'fe_static':
            time_series_attribute_columns.append(name)
          elif details['tag'] == 'fe_available_past_future':
            available_at_forecast_columns.append(name)
          elif details['tag'] == 'fe_available_past_only':
            unavailable_at_forecast_columns.append(name)

          # Determine data type of the column. Note: columns without transformations
          # (except for target) will be ignored during training.
          if feature_columns and name not in feature_columns:
            continue
          if details['type'] == 'STRING':
            trans = {'categorical': {'column_name': name}}
            column_transformations.append(trans)
          elif details['type'] in ['INT64', 'INTEGER', 'NUMERIC', 'FLOAT64', 'FLOAT']:
            trans = {'numeric': {'column_name': name}}
            column_transformations.append(trans)
          elif details['type'] in ['DATETIME', 'DATE', 'TIMESTAMP']:
            trans = {'timestamp': {'column_name': name}}
            column_transformations.append(trans)

        return NamedTuple('Outputs', [
            ('time_series_identifier_column', str),
            ('time_series_attribute_columns', list),
            ('available_at_forecast_columns', list),
            ('unavailable_at_forecast_columns', list),
            ('column_transformations', list),
            ('preprocess_bq_uri', str),
            ('target_column', str),
            ('time_column', str),
            ('predefined_split_column', str),
            ('weight_column', str),
            ('data_granularity_unit', str),
            ('data_granularity_count', int),
        ])(
            time_series_identifier_column,
            time_series_attribute_columns,
            available_at_forecast_columns,
            unavailable_at_forecast_columns,
            column_transformations,
            bigquery_table_uri,
            primary_metadata['target_column'],
            primary_metadata['time_column'],
            predefined_split_column,
            weight_column,
            primary_metadata['time_granularity']['unit'].lower(),
            primary_metadata['time_granularity']['quantity'],
        )

      def _serialize_int(int_value: int) -> str:
          if isinstance(int_value, str):
              return int_value
          if not isinstance(int_value, int):
              raise TypeError('Value "{}" has type "{}" instead of int.'.format(
                  str(int_value), str(type(int_value))))
          return str(int_value)

      def _serialize_json(obj) -> str:
          if isinstance(obj, str):
              return obj
          import json

          def default_serializer(obj):
              if hasattr(obj, 'to_struct'):
                  return obj.to_struct()
              else:
                  raise TypeError(
                      "Object of type '%s' is not JSON serializable and does not have .to_struct() method."
                      % obj.__class__.__name__)

          return json.dumps(obj, default=default_serializer, sort_keys=True)

      def _serialize_str(str_value: str) -> str:
          if not isinstance(str_value, str):
              raise TypeError('Value "{}" has type "{}" instead of str.'.format(
                  str(str_value), str(type(str_value))))
          return str_value

      import json
      import argparse
      _parser = argparse.ArgumentParser(prog='Prepare data for train', description='Prepares the parameters for the training step.')
      _parser.add_argument("--input-tables", dest="input_tables", type=json.loads, required=True, default=argparse.SUPPRESS)
      _parser.add_argument("--preprocess-metadata", dest="preprocess_metadata", type=json.loads, required=True, default=argparse.SUPPRESS)
      _parser.add_argument("--model-feature-columns", dest="model_feature_columns", type=json.loads, required=False, default=argparse.SUPPRESS)
      _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=12)
      _parsed_args = vars(_parser.parse_args())
      _output_files = _parsed_args.pop("_output_paths", [])

      _outputs = prepare_data_for_train(**_parsed_args)

      _output_serializers = [
          _serialize_str,
          _serialize_json,
          _serialize_json,
          _serialize_json,
          _serialize_json,
          _serialize_str,
          _serialize_str,
          _serialize_str,
          _serialize_str,
          _serialize_str,
          _serialize_str,
          _serialize_int,

      ]

      import os
      for idx, output_file in enumerate(_output_files):
          try:
              os.makedirs(os.path.dirname(output_file))
          except OSError:
              pass
          with open(output_file, 'w') as f:
              f.write(_output_serializers[idx](_outputs[idx]))
    args:
    - --input-tables
    - {inputValue: input_tables}
    - --preprocess-metadata
    - {inputValue: preprocess_metadata}
    - if:
        cond: {isPresent: model_feature_columns}
        then:
        - --model-feature-columns
        - {inputValue: model_feature_columns}
    - '----output-paths'
    - {outputPath: time_series_identifier_column}
    - {outputPath: time_series_attribute_columns}
    - {outputPath: available_at_forecast_columns}
    - {outputPath: unavailable_at_forecast_columns}
    - {outputPath: column_transformations}
    - {outputPath: preprocess_bq_uri}
    - {outputPath: target_column}
    - {outputPath: time_column}
    - {outputPath: predefined_split_column}
    - {outputPath: weight_column}
    - {outputPath: data_granularity_unit}
    - {outputPath: data_granularity_count}
