%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Copyright (c) 2023, Amon Lahr, Simon Muntwiler, Antoine Leeman & Fabian Flürenbrock Institute for Dynamic Systems and Control, ETH Zurich.
%
% All rights reserved.
%
% Please see the LICENSE file that has been included as part of this package.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [tuning_struct, i_opt] = lqr_tuning(x0,Q,params)
    % YOUR CODE HERE
    i_opt = nan;
    best_J_u = inf;

    % Prepare an array of empty structs
    tuning_struct = repmat(struct( ...
        'InitialCondition', {}, ...
        'Qdiag', {}, ...
        'MaxAbsPositionXZ', {}, ...
        'MaxAbsPositionY', {}, ...
        'MaxAbsThrust', {}, ...
        'InputCost', {}, ...
        'MaxFinalPosDiff', {}, ...
        'MaxFinalVelDiff', {}, ...
        'TrajFeasible', {} ...
    ), size(Q,2), 1);

    for i=1:size(Q,2)
        tuning_struct(i).InitialCondition = x0;
        tuning_struct(i).Qdiag = Q(:,i);

        ctrl = LQR(diag(Q(:,i)), eye(params.model.nu), params);
        [Xt, Ut, ~] = simulate(x0, ctrl, params);

        [s_max, y_max, u_max, J_u, df_max, vf_max, traj_feas] = ...
            traj_constraints(Xt, Ut, params);

        tuning_struct(i).MaxAbsPositionXZ = s_max;
        tuning_struct(i).MaxAbsPositionY = y_max;
        tuning_struct(i).MaxAbsThrust = u_max;
        tuning_struct(i).InputCost = J_u;
        tuning_struct(i).MaxFinalPosDiff = df_max;
        tuning_struct(i).MaxFinalVelDiff = vf_max;
        tuning_struct(i).TrajFeasible = traj_feas;

        if traj_feas
            if J_u < best_J_u
                i_opt = i;
                best_J_u = J_u;
            end
        end
    end

    % because the test suite wants a column vector
    tuning_struct = tuning_struct';
end