client = jcmoptimizer.Client(); 

%Rastrigin-like function depending on additional phase offset phi 
function val=rast(x1, x2, phi)
    val = 10*2 + x1^2 + x2^2 - 10*cos(2*pi*x1 + phi) - 10*cos(2*pi*x2);
end
 
%time-dependent slowly varying phi
function phi=current_phi()
    t = now()*1e5;
    phi = 2*pi*sin(t/180);
end
 
% Definition of the search domain
design_space = { ...
    struct('name', 'x1', 'type', 'continuous', 'domain', [-1.5,1.5]), ...
    struct('name', 'x2', 'type', 'continuous', 'domain', [-1.5,1.5]) ...
};

% Definition of the environment variable "phi"
environment = { ...
    struct('name', 'phi', 'type', 'variable', 'domain', [-2*pi, 2*pi]) ...
};

 % Creation of the study object with study_id 'changing_environment'
study = client.create_study( ...
    'design_space', design_space, ...
    'environment', environment, ...
    'driver','ActiveLearning',...
    'study_name','Optimal control of a system in a changing environment',...
    'study_id', 'changing_environment');

% In the initial training phase, the target is to explore the
% parameter space to find the global minimim.
study.configure( ...
    ... train with 500 data points
    'max_iter', 500, ...    
    ... Advanced sample computation is switched off since the environment
    ... parameter phi can change significantly between computation
    ... of the suggestion and evaluation of the objective function
    'acquisition_optimizer', struct('compute_suggestion_in_advance', false) ...
);

% Evaluation of the black-box function for specified design parameters
function observation = evaluate(study, sample)

    pause(2); % make objective expensive
    observation = study.new_observation();
    % get current phi
    phi = current_phi();
    observation.add(rast(sample.x1, sample.x2, phi), 'environment_value', {phi});
    
end  

% Run the minimization
study.set_evaluator(@evaluate);
study.run(); 


% The target in the control phase is to evaluate the offet Rastrigin function only
% at well performing (x1,x2)-point depending on the current value of the environment.
MAX_ITER = 500; %evaluate for 500 additional iterations
study.configure( ...
    'max_iter', 500 + MAX_ITER, ...
    ... The scaling is reduced to penalize parameters with large uncertainty    
    'scaling', 0.01, ...
    ... The lower-confidence bound (LCB) strategy is chosen instead of the
    ... default expected improvement (EI). LCB is easier to maximize at the
    ... risk of less exploration of the parameter space, which is anyhow not
    ... desired in the control phase.
    'objectives', {...
        struct('type', 'Minimizer', 'name', 'objective', 'strategy', 'LCB') ...
    }, ...
    'acquisition_optimizer', struct('compute_suggestion_in_advance', false) ...
);

% keep track of suggested design points and phis at request time and evaluation time
design_points = {};
phis_at_request = {};
phis_at_eval = {};
            
iter = 0;
while not(study.is_done())
    iter = iter + 1;
    if iter > MAX_ITER
        break
    end
        
    phi = current_phi();
    suggestion = study.get_suggestion({phi});
    phis_at_request{end + 1} = phi;
    sample = suggestion.sample;
    design_points{end + 1} = [sample.x1, sample.x2]; 
    try
        obs = evaluate(study, sample);
        %update phi from observation
        phi = obs.data.None{1}.env{1};
        phis_at_eval{end + 1} = phi;
        
        predictions = study.driver.predict({[sample.x1, sample.x2, phi]});
        std = sqrt(predictions.variance(1, 1));
            
        fprintf("Uncertainty of prediction %f\n", std);    
        % add data only if prediction has significant uncertainty
        if std > 0.01
            study.add_observation(obs, suggestion.id);
        else
            study.clear_suggestion( ...
                suggestion.id, sprintf("Ignoring observation with uncertainty %f", std));
        end
            
    catch err
        study.clear_suggestion( ...
            suggestion.id, ...
            sprintf("Evaluator function failed with error: %s", err.message) ...
        )
        rethrow(err);
    end    

end
            

fig = figure('Position', [0, 0, 1000, 500]);
 
% all observed training samples
observed = study.driver.get_observed_values();
subplot(1, 2, 1)
plot(observed.means,".")
hold on;
xline(500, '--');
xlabel("training+control iteration")
ylabel("observed value of Rastrigin function")

%observed values during control phase
observed_vals = [];

%values that would have been observed at request time,
%i.e. if there would be no time delay between request and 
%evaluation of suggestion
observed_vals_at_request = [];

for i = 1:length(design_points)
    observed_vals(i) = rast( ...
        design_points{i}(1), design_points{i}(2), phis_at_eval{i});
    observed_vals_at_request(i) = rast( ...
        design_points{i}(1), design_points{i}(2), phis_at_request{i});
end

%best value of x1-parameter depending on environment
function val = best_x1(phi)
    val = -phi/(2*pi);
    if abs(phi) > pi
        val = val + sign(phi);
    end
end

% best possible values 
best_vals = [];
for i = 1:length(phis_at_eval)
    phi = phis_at_eval{i};
    best_vals(i) = rast(best_x1(phi), 0, phi);
end

subplot(1, 2, 2)
semilogy(observed_vals,".", 'DisplayName', 'observed values')
hold on;
semilogy(observed_vals_at_request,".", 'DisplayName', 'observed values if no time delay')
semilogy(best_vals, 'DisplayName', 'best possible values')
ylim([1e-4, 1e1])
xlabel("control iteration")
legend()
saveas(fig, "training_and_control.png")

client.shutdown_server();