# Copyright 2021 The Kubeflow Authors
#
# Licensed 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.

name: training_configurator_and_validator
description: |
  Component to configure training and validate data and user-input configurations.

  Args:
      dataset_stats (AutoMLTabularDatasetStats):
          Dataset stats generated by feature transform engine.
      split_example_counts (str):
          JSON string of data split example counts for train, validate, and test splits.
      training_schema_path (DatasetSchema):
          Schema of input data to the tf_model at training time.
      instance_schema (InstanceSchema):
          Schema of input data to the tf_model at serving time.
      target_column (str):
          Target column of input data.
      weight_column (str):
          Weight column of input data.
      prediction_type (str):
          Model prediction type. One of "classification", "regression", "time_series".
      optimization_objective (str):
          Objective function the model is optimizing towards. The training process creates a model
          that maximizes/minimizes the value of the objective function over the validation set.
          The supported optimization objectives depend on the prediction type. If the field is not
          set, a default objective function is used.
          classification (binary):
              "maximize-au-roc" (default) - Maximize the area under the receiver
                  operating characteristic (ROC) curve.
              "minimize-log-loss" - Minimize log loss.
              "maximize-au-prc" - Maximize the area under the precision-recall curve.
              "maximize-precision-at-recall" - Maximize precision for a specified recall value.
              "maximize-recall-at-precision" - Maximize recall for a specified precision value.
          classification (multi-class):
              "minimize-log-loss" (default) - Minimize log loss.
          regression:
              "minimize-rmse" (default) - Minimize root-mean-squared error (RMSE).
              "minimize-mae" - Minimize mean-absolute error (MAE).
              "minimize-rmsle" - Minimize root-mean-squared log error (RMSLE).
      optimization_objective_recall_value (str):
          Required when optimization_objective is "maximize-precision-at-recall". Must be between
          0 and 1, inclusive.
      optimization_objective_precision_value (str):
          Required when optimization_objective is "maximize-recall-at-precision". Must be between
          0 and 1, inclusive.
      run_evaluation (bool):
        Whether we are running evaluation in the training pipeline.
      run_distill (bool):
          Whether the distillation should be applied to the training.
      enable_probabilistic_inference (bool): If probabilistic inference is enabled, the
        model will fit a distribution that captures the uncertainty of a
        prediction. At inference time, the predictive distribution is used to make
        a point prediction that minimizes the optimization objective. For example,
        the mean of a predictive distribution is the point prediction that
        minimizes RMSE loss. If quantiles are specified, then the quantiles of the
        distribution are also returned.
      time_series_identifier_column (str):
        Time series idenfier column. Used by forecasting only.
      time_column (str):
        The column that indicates the time. Used by forecasting only.
      time_series_attribute_columns (str):
        The column names of the time series attributes.
      available_at_forecast_columns (str):
        The names of the columns that are available at forecast time.
      unavailable_at_forecast_columns (str):
        The names of the columns that are not available at forecast time.
      quantiles (str):
        All quantiles that the model need to predict.
      context_window (int):
        The length of the context window.
      forecast_horizon (int):
        The length of the forecast horizon.
      forecasting_model_type (str):
        The model types, e.g. l2l, seq2seq, tft.
      forecasting_transformations_path (str):
        The path to the JSON format forecasting transformations. Used by forecasting only.
      stage_1_deadline_hours (Optional[float]):
        Stage 1 training budget in hours.
      stage_2_deadline_hours (Optional[float]):
        Stage 2 training budget in hours.

  Returns:
      metadata (TabularExampleGenMetadata):
          The tabular example gen metadata.

inputs:
- {name: dataset_stats, type: AutoMLTabularDatasetStats}
- {name: split_example_counts, type: String}
- {name: training_schema, type: DatasetSchema}
- {name: instance_schema, type: DatasetSchema}
- {name: target_column, type: String, default: ""}
- {name: weight_column, type: String, default: ""}
- {name: prediction_type, type: String, default: ""}
- {name: optimization_objective, type: String, default: ""}
- {name: optimization_objective_recall_value, type: Float, default: "-1"}
- {name: optimization_objective_precision_value, type: Float, default: "-1"}
- {name: run_evaluation, type: Boolean, default: "false"}
- {name: run_distill, type: Boolean, default: "false"}
- {name: enable_probabilistic_inference, type: Boolean, default: "false"}
- {name: time_series_identifier_column, type: String, default: ""}
- {name: time_column, type: String, default: ""}
- {name: time_series_attribute_columns, type: String, default: ""}
- {name: available_at_forecast_columns, type: String, default: ""}
- {name: unavailable_at_forecast_columns, type: String, default: ""}
- {name: quantiles, type: String, default: ""}
- {name: context_window, type: Integer, default: "-1"}
- {name: forecast_horizon, type: Integer, default: "-1"}
- {name: forecasting_model_type, type: String, default: ""}
- {name: forecasting_transformations_path, type: String, default: ""}
- {name: stage_1_deadline_hours, type: Float, optional: true}
- {name: stage_2_deadline_hours, type: Float, optional: true}

outputs:
- {name: metadata, type: TabularExampleGenMetadata}
- {name: instance_baseline, type: AutoMLTabularInstanceBaseline}

implementation:
  container:
    image: us-docker.pkg.dev/vertex-ai/automl-tabular/feature-transform-engine:20230322_0325
    args: [
      training_configurator_and_validator,
      concat: [
        "--instance_schema_path=",
        {inputUri: instance_schema}
      ],
      concat: [
        "--training_schema_path=",
        {inputUri: training_schema}
      ],
      concat: [
        "--dataset_stats_path=",
        {inputUri: dataset_stats}
      ],
      concat: [
        "--split_example_counts=",
        {inputValue: split_example_counts}
      ],
      concat: [
        "--target_column=",
        {inputValue: target_column}
      ],
      concat: [
        "--weight_column=",
        {inputValue: weight_column}
      ],
      concat: [
        "--prediction_type=",
        {inputValue: prediction_type}
      ],
      concat: [
        "--optimization_objective=",
        {inputValue: optimization_objective}
      ],
      concat: [
        "--optimization_objective_recall_value=",
        {inputValue: optimization_objective_recall_value}
      ],
      concat: [
        "--optimization_objective_precision_value=",
        {inputValue: optimization_objective_precision_value}
      ],
      concat: [
        "--metadata_path=",
        {outputUri: metadata}
      ],
      concat: [
        "--instance_baseline_path=",
        {outputUri: instance_baseline}
      ],
      concat: [
        "--run_evaluation=",
        {inputValue: run_evaluation}
      ],
      concat: [
        "--run_distill=",
        {inputValue: run_distill}
      ],
      concat: [
        "--enable_probabilistic_inference=",
        {inputValue: enable_probabilistic_inference}
      ],
      concat: [
        "--time_series_identifier_column=",
        {inputValue: time_series_identifier_column}
      ],
      concat: [
        "--time_column=",
        {inputValue: time_column}
      ],
      concat: [
        "--time_series_attribute_columns=",
        {inputValue: time_series_attribute_columns}
      ],
      concat: [
        "--available_at_forecast_columns=",
        {inputValue: available_at_forecast_columns}
      ],
      concat: [
        "--unavailable_at_forecast_columns=",
        {inputValue: unavailable_at_forecast_columns}
      ],
      concat: [
        "--quantiles=",
        {inputValue: quantiles}
      ],
      concat: [
        "--context_window=",
        {inputValue: context_window}
      ],
      concat: [
        "--forecast_horizon=",
        {inputValue: forecast_horizon}
      ],
      concat: [
        "--forecasting_model_type=",
        {inputValue: forecasting_model_type}
      ],
      concat: [
        "--forecasting_transformations_path=",
        {inputValue: forecasting_transformations_path}
      ],
      {if: {
        cond: {isPresent: stage_1_deadline_hours},
        then:
          {concat: [
            "--stage_1_deadline_hours=",
            {inputValue: stage_1_deadline_hours}
          ]},
      }},
      {if: {
        cond: {isPresent: stage_2_deadline_hours},
        then:
          {concat: [
            "--stage_2_deadline_hours=",
            {inputValue: stage_2_deadline_hours}
          ]},
      }}
    ]
