# This pipeline demonstrates hyper-parameter optimization.
# The goal is to find a set of hyper-parameter values that helps train the best model.
# We launch several optimization stages sequentially.
# At each stage the optimizer suggests several parameter sets to explore based on the available measurements.
# For each suggested parameter set we train a model (semi-dummy) and measure its quality metrics.
# We then collect the metrics for all suggested parameter sets and update out measurements set.
# With the expanded set of measurements, each new optimization stage should result in better parameter set suggestions.
#
# One aspect of this pipeline is the atomicity of the parameter set suggestion.
# Some optimizers have a persistent mutable global state that is changed when parameter set metrics are submitted.
# The presence of mutable global state may cause reproducibility issues where suggestions for a new model might be based on measurements from a different model.
# The "suggest_parameter_sets_from_measurements_op" in this pipeline is a single operation, which behaves like a pure function and does not rely on external global state.

kfp_endpoint = None


import kfp
from kfp import components


suggest_parameter_sets_from_measurements_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/382c4d109fbd489bd85de54dd9171150e326b401/components/google-cloud/Optimizer/Suggest_parameter_sets_based_on_measurements/component.yaml')

get_element_by_index_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/55ef28a9d51edc4eeed2a5c6f44cc7457e8a41d8/components/json/Get_element_by_index/component.yaml')
build_dict_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/4a4be6b748b0d1284d65a417ce4ab5bec596e9fe/components/json/Build_dict/component.yaml')
build_list_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/4a4be6b748b0d1284d65a417ce4ab5bec596e9fe/components/json/Build_list/component.yaml')
combine_lists_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/4a4be6b748b0d1284d65a417ce4ab5bec596e9fe/components/json/Combine_lists/component.yaml')


# The train_and_measure_model is a semi-dummy component that creates a model given the [hyper]parameters and evaluates that model.
# In this case, the model is a polinomial model.
# The evaluation procedure compares the model with the real function that our model is trying to learn
# and calculates the mean squared error based on a random sample of data points.
# In real world cases this component will be substituted by a sequence of model trainer, predictor and evaluator components.
@components.create_component_from_func
def train_and_measure_model(parameters: dict) -> float:
    import random

    def real_function(x):
        p1 = 3
        p2 = -1
        p3 = 2
        return p1 * x**2 + p2 * x + p3

    def get_eval_set() -> dict:
        eval_set = {}
        num_samples = 100
        for i in range(num_samples):
            x = random.normalvariate(0, 1) * 5
            eval_set[x] = real_function(x)
        return eval_set

    def train_model(parameters):
        def apply_model(x):
            return parameters['p1'] * x**2 + parameters['p2'] * x + parameters['p3']
        return apply_model

    model = train_model(parameters)

    eval_set = get_eval_set()
    sum_squared_error = 0
    
    for x, expected_y in eval_set.items():
        actual_y = model(x)
        error = abs(expected_y - actual_y)
        squared_error = error ** 2
        sum_squared_error += squared_error
    mean_squared_error = sum_squared_error / len(eval_set)
    return mean_squared_error


parameter_specs=[
    {
        'parameter': 'p1',
        'type': 'DOUBLE',
        'double_value_spec' : {
            'min_value' : -5,
            'max_value' : 5,
        }
    },
    {
        'parameter': 'p2',
        'type': 'DOUBLE',
        'double_value_spec': {
            'min_value': -5,
            'max_value': 5,
        }
    },
    {
        'parameter': 'p3',
        'type': 'DOUBLE',
        'double_value_spec': {
            'min_value': -5,
            'max_value': 5,
        }
    },
]


def optimizer_pipeline():
    # Number of optimization stages and suggestions per stage.
    # Note that these numbers cannot be parametrized, since they're used in compile-time python loops.
    optimization_stages = 3
    suggestions_per_stage = 5

    # We launch several optimization stages sequentially.
    # At each stage the optimizer suggests several parameter sets to explore based on the available measurements.
    # Each stage depends on the completion of all trials in the previous stage (since only completed trials affect new trial suggesions).
    # Each optimization stage should result in better parameter set suggestions.
    all_metrics_for_parameter_sets = []
    for stage in range(optimization_stages):
        parameter_sets = suggest_parameter_sets_from_measurements_op(
            parameter_specs=parameter_specs,
            metrics_for_parameter_sets=all_metrics_for_parameter_sets,
            suggestion_count=suggestions_per_stage,
            maximize=False,
        ).output

        # Evaluate each suggested set of parameters.
        # Loop over the suggested trials.
        # We need to collect the created tasks in the `trial_measurement_tasks` list so that the next round of suggestions can depend on their completion.
        # Cannot use dsl.ParallelFor here due to a bug in Argo https://github.com/argoproj/argo-workflows/issues/2660
        # Without ParallelFor we have to use python loop
        # and explicitly get individual suggections using the get_element_by_index_op component
        # then extract the trial name and parameter sets using get_element_by_key_op and query_json_op components.
        new_metrics_for_parameter_sets = []
        for siggestion_index in range(suggestions_per_stage):
            parameter_set = get_element_by_index_op(
                json=parameter_sets,
                index=siggestion_index,
            ).output

            model_error = train_and_measure_model(
                parameters=parameter_set,
            ).output

            metric_for_parameter_set = build_dict_op(
                key_1='parameters',
                value_1=parameter_set,
                key_2='metrics',
                value_2={
                    'metric': model_error,
                },
            ).output

            new_metrics_for_parameter_sets.append(metric_for_parameter_set)
        # Collecting metrics for the current stage
        new_list_of_metrics_for_parameter_sets = build_list_op(*new_metrics_for_parameter_sets).output
        # Collecting metrics for all stages
        all_metrics_for_parameter_sets = combine_lists_op(all_metrics_for_parameter_sets, new_list_of_metrics_for_parameter_sets).output


if __name__ == '__main__':
    kfp.Client(host=kfp_endpoint).create_run_from_pipeline_func(optimizer_pipeline, arguments={})
