%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 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
    R = eye(params.model.nu);
    size_Q = size(Q);
    param_choices = size_Q(2);
    J_array = [];
    tuning_struct = [];
    for i = 1:param_choices
        Q_mat = diag(Q(:,i));
        ctrl = LQR(Q_mat,R,params);
        [Xt, Ut, ctrl_info] = simulate(x0, ctrl, params);
        ctrl_feas_arr = [];
        for k = 1:params.model.HorizonLength
            ctrl_feas_arr = [ctrl_feas_arr,ctrl_info(k).ctrl_feas];
        end
        if all(ctrl_feas_arr)
            J = 0;
            for j = 1:params.model.HorizonLength
                J = J + Xt(:,j)'*Q_mat*Xt(:,j);
                J = J + Ut(:,j)'*R*Ut(:,j);
            end
            J = J + Xt(:,end)'*Q_mat*Xt(:,end);
            J_array = [J_array, J];
            tuning_struct_item.TrajFeasible = true;
        else
            tuning_struct_item.TrajFeasible = false;
        end
        tuning_struct_item.InitialCondition = x0;
        tuning_struct_item.Qdiag = Q(:,i);
        tuning_struct_item.MaxAbsPositionXZ = params.constraints.MaxAbsPositionXZ;
        tuning_struct_item.MaxAbsPositionY = params.constraints.MaxAbsPositionY;
        tuning_struct_item.MaxAbsThrust = params.constraints.MaxAbsThrust;
        tuning_struct_item.InputCost = J_array(i);
        tuning_struct_item.MaxFinalPosDiff = params.constraints.MaxFinalPosDiff;
        tuning_struct_item.MaxFinalVelDiff = params.constraints.MaxFinalVelDiff;
        tuning_struct = [tuning_struct,tuning_struct_item];
    end
%     display(J_array)
    tuning_struct = tuning_struct'; 
    if isempty(J_array)
        i_opt = nan;
    else
        [~, i_opt] = min(J_array);
    end
end