classdef quadratically_constrained_quadratic_bilevel
    % Quadratically Constrained Quadratic Bilevel Program (QCQBP)
    % ===========================================================
    % The objective function and constraints are quadratic expression in (x, y).
    % Each quadratic expression is parameterised by (c, d, P, R, Q, s).
    %
    %
    % Quadratic expression in x and y
    % ===============================
    %     x: Leader's decision variables; vector of shape (n_x,).
    %     y: Follower's decision variables; vector of shape (n_y,).
    %     c: Linear coefficient of x. vector of shape (n_x,).
    %     d: Linear coefficient of y. shape (n_y,).
    %     P: Quadratic form matrix of shape (n_x, n_x).
    %     Q: Quadratic form matrix of shape (n_y, n_y).
    %     R: Quadratic form matrix of shape (n_x, n_y).
    %     s: Scalar constant.
    %
    %
    %     Q(x,y) := 0.5 [x]'* [ P  R ] * [x]    + [c]'[x]    + s
    %                   [y]   [ R' Q ]   [y]      [d] [y]
    %


    properties(Constant)
        name = 'quadratically_constrained_quadratic_bilevel';
        category = 'archetype';
        subcategory = '';
        author = 'Samuel Ward';
        datasets = {'allende_still_2013.json', 'bard_1988_ex3.json', 'dempe_1992a.json', 'dempe_1992b.json'};
        paths = fullfile('bolib3', 'data', 'quadratically_constrained_quadratic_bilevel', quadratic_bilevel.datasets);
    end


    methods(Static)

        function val = quadratic(x, y, coefficients)
            % Quadratic in
            %     x: leader's decision variables; vector of shape (n_x,).
            %     y: follower's decision variables; vector of shape (n_y,).
            %
            % The coefficients argument should be a stuct with feilds
            %     c: linear coefficient of x. vector of shape (n_x,).
            %     d: linear coefficient of y. shape (n_y,).
            %     P: quadratic form matrix of shape (n_x, n_x).
            %     Q: quadratic form matrix of shape (n_y, n_y).
            %     R: quadratic form matrix of shape (n_x, n_y).
            %     s: scalar constant.
            val = 0.5 * x' * coefficients.P * x + ...
                  0.5 * y' * coefficients.Q * y + ...
                  1.0 * x' * coefficients.R * y + ...
                  dot(coefficients.c, x) + ...
                  dot(coefficients.d, y) + ...
                  coefficients.s;
        end

        function val = F(x, y, data)
            % Upper-level objective function
            val = quadratically_constrained_quadratic_bilevel.quadratic( ...
                x, y, data.F);
        end

        function ineq = G(x, y, data)
            % Upper-level inequality constraints
            ineq = [];
            if ~isempty(data.ineq_upper)
                ineq = zeros(length(data.ineq_upper), 1);
                for i = 1:length(data.ineq_upper)
                    ineq(i) = quadratically_constrained_quadratic_bilevel.quadratic( ...
                        x, y, data.ineq_upper(i));
                end
            end
            if ~isempty(data.rhs_upper)
                ineq = [ineq; data.A_upper * x + data.B_upper * y - data.rhs_upper];
            end
            if ~isempty(data.x_lower_bound)
                ineq = [ineq; x - data.x_lower_bound];
            end
            if ~isempty(data.x_upper_bound)
                ineq = [ineq; data.x_upper_bound - x];
            end
        end

        function val = H(~, ~, ~)
            % Upper-level equality constraints
            val = [];
        end

        function val = f(x, y, data)
            % Lower-level objective function
            val = quadratically_constrained_quadratic_bilevel.quadratic( ...
                x, y, data.f);
        end

        function val = g(x, y, data)
            % Lower-level inequality constraints
            val = [];
            if ~isempty(data.ineq_lower)
                val = zeros(length(data.ineq_lower), 1);
                for i = 1:length(data.ineq_lower)
                    val(i) = quadratically_constrained_quadratic_bilevel.quadratic( ...
                        x, y, data.ineq_lower(i));
                end
            end
            if ~isempty(data.rhs_lower)
                val = [val; data.A_lower * x + data.B_lower * y - data.rhs_lower];
            end
            if ~isempty(data.y_lower_bound)
                val = [val; y - data.y_lower_bound];
            end
            if ~isempty(data.y_upper_bound)
                val = [val; data.y_upper_bound - y];
            end
        end

        function val = h(~, ~, ~)
            % Lower-level equality constraints
            val = [];
        end


        % %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% %
        % -=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=- %
        %%                         Derivatives                         %%
        % -=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=- %
        % %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% %
        function dx = dF_dx(x, y, data)
            % ∂F/∂x
            % Dimention (n_x × 1)
            dx = (data.F.P * x) + (data.F.R * y) + data.F.c;
        end

        function dy = dF_dy(x, y, data)
            % ∂F/∂y
            % Dimention (n_y × 1)
            dy = (data.F.Q * y) + (data.F.R' * x) + data.F.d;
        end

        function Hxx = d2F_dxx(~, ~, data)
            % ∂²F/∂x²
            % Dimention (n_x × n_x)
            Hxx = data.F.P;
        end

        function Hyy = d2F_dyy(~, ~, data)
            % ∂²F/∂y²
            % Dimention (n_y × n_y)
            Hyy = data.F.Q;
        end

        function Hxy = d2F_dxy(~, ~, data)
            % ∂²F/∂x∂y
            % Dimention (n_x × n_y)
            Hxy = data.F.R;
        end

        function Jx = dG_dx(x, y, data)
            % Jacobian of G wrt x
            % Dimention (|G| × n_x)
            Jx = zeros(0, data.n_x);
            if ~isempty(data.ineq_upper)
                Jx = zeros(length(data.ineq_upper), data.n_x);
                for i = 1:length(data.ineq_upper)
                    ineq = data.ineq_upper(i);
                    Jx(i,:) = (ineq.P * x) + (ineq.R * y) + ineq.c;
                end
            end
            if ~isempty(data.rhs_upper)
                Jx = [Jx; data.A_upper];
            end
            if ~isempty(data.x_lower_bound)
                Jx = [Jx; eye(data.n_x)];
            end
            if ~isempty(data.x_upper_bound)
                Jx = [Jx; -eye(data.n_x)];
            end
        end

        function Jy = dG_dy(x, y, data)
            % Jacobian of G wrt y:
            % Dimention (|G| × n_y)
            Jy = zeros(0, data.n_y);
            if ~isempty(data.ineq_upper)
                Jy = zeros(length(data.ineq_upper), data.n_x);
                for i = 1:length(data.ineq_upper)
                    ineq = data.ineq_upper(i);
                    Jy(i,:) = (ineq.Q * y) + (ineq.R' * x) + ineq.d;
                end
            end
            if ~isempty(data.rhs_upper)
                Jy = [Jy; data.B_upper];
            end
            if ~isempty(data.x_lower_bound)
                Jy = [Jy; zeros(data.n_x, data.n_y)];
            end
            if ~isempty(data.x_upper_bound)
                Jy = [Jy; zeros(data.n_x, data.n_y)];
            end
        end

        function Jx = dH_dx(~, ~, data)
            % Jacobian of H wrt x
            % Dimention (0 × n_x)
            Jx = zeros(0, data.n_x);
        end

        function Jy = dH_dy(~, ~, data)
            % Jacobian of H wrt y
            % Dimention (0 × n_y)
            Jy = zeros(0, data.n_y);
        end

        function dx = df_dx(x, y, data)
            % ∂f/∂x
            % Dimention (n_x × 1)
            dx = (data.f.P * x) + (data.f.R * y) + data.f.c;
        end

        function dy = df_dy(x, y, data)
            % ∂f/∂y
            % Dimention (n_y × 1)
            dy = (data.f.Q * y) + (data.f.R' * x) + data.f.d;
        end

        function Hxx = d2f_dxx(~, ~, data)
            % ∂²f/∂x²
            % Dimention (n_x × n_x)
            Hxx = data.f.P;
        end

        function Hyy = d2f_dyy(~, ~, data)
            % ∂²f/∂y²
            % Dimention (n_y × n_y)
            Hyy = data.f.Q;
        end

        function Hxy = d2f_dxy(~, ~, data)
            % ∂²f/∂x∂y
            % Dimention (n_x × n_y)
            Hxy = data.f.R;
        end

        function Jx = dg_dx(x, y, data)
            % Jacobian of g wrt x
            % Dimention (|g| × n_x)
            Jx = zeros(0, data.n_x);
            if ~isempty(data.ineq_lower)
                Jx = zeros(length(data.ineq_lower), data.n_x);
                for i = 1:length(data.ineq_lower)
                    ineq = data.ineq_lower(i);
                    Jx(i,:) = (ineq.P * x) + (ineq.R * y) + ineq.c;
                end
            end
            if ~isempty(data.rhs_lower)
                Jx = [Jx; data.A_lower];
            end
            if ~isempty(data.y_lower_bound)
                Jx = [Jx; zeros(data.n_y, data.n_x)];
            end
            if ~isempty(data.y_upper_bound)
                Jx = [Jx; zeros(data.n_y, data.n_x)];
            end
        end

        function Jy = dg_dy(x, y, data)
            % Jacobian of g wrt y
            % Dimention (|g| × n_y)
            Jy = zeros(0, data.n_y);
            if ~isempty(data.ineq_lower)
                Jy = zeros(length(data.ineq_lower), data.n_x);
                for i = 1:length(data.ineq_lower)
                    ineq = data.ineq_lower(i);
                    Jy(i,:) = (ineq.Q * y) + (ineq.R' * x) + ineq.d;
                end
            end
            if ~isempty(data.rhs_lower)
                Jy = [Jy; data.B_lower];
            end
            if ~isempty(data.y_lower_bound)
                Jy = [Jy; eye(data.n_y)];
            end
            if ~isempty(data.y_upper_bound)
                Jy = [Jy; -eye(data.n_y)];
            end
        end

        function Jx = dh_dx(~, ~, data)
            % Jacobian of h wrt x
            % Dimention (0 × n_x)
            Jx = zeros(0, data.n_x);
        end

        function Jy = dh_dy(~, ~, data)
            % Jacobian of h wrt y
            % Dimention (0 × n_y)
            Jy = zeros(0, data.n_y);
        end


        % %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% %
        % -=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=- %
        %%                          Read Data                          %%
        % -=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=- %
        % %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% %
        function data = read_data(path)
            % Read JSON file
            txt = fileread(path);
            data = jsondecode(txt);

            % Store the path. This is useful diagnostics
            data.path = char(path);

            % ================== REQUIRED keys ==================
            % Objective coefficients are nessasary
            % Convert the coefficents to double data type
            keys = {'F', 'f'};
            subkeys = {'c', 'd', 'P', 'Q', 'R', 's'};
            for i = 1:length(keys)
                key = keys{i};
                % Diagnostics message
                msg = sprintf(strcat('Invalid quadratic form for %s. ', ...
                    'It should look like this ', ...
                    '\n    "%s": { "c": [], "d": [], "P": [[]], "Q": [[]], "R": [[]], "s": 0 } ', ...
                    '\n    See JSON file: %s.'), key, key, path ...
                );
                assert(isfield(data, key), msg)
                for j = 1:length(subkeys)
                    subkey = subkeys{j};
                    assert(isfield(data.(key), subkey), msg)
                    data.(key).(subkey) = double(data.(key).(subkey));
                end
            end


            % ================== OPTIONAL keys ==================
            % Bounds are linear ineq constraints are optional
            % They should be cast to the double data type
            % Or default to the correct empty shape
            keys = {
                "A_upper"; "B_upper"; "c_upper"; "d_upper"; "rhs_upper"; "x_upper_bound";  "x_lower_bound"; ...
                "A_lower"; "B_lower"; "c_lower"; "d_lower"; "rhs_lower"; "y_upper_bound";  "y_lower_bound";...
            };
            default_values = {
                zeros(0, data.n_x); zeros(0, data.n_y); zeros(0, 1); zeros(0, 1); zeros(0, 1); zeros(0, 1); zeros(0, 1); ...
                zeros(0, data.n_x); zeros(0, data.n_y); zeros(0, 1); zeros(0, 1); zeros(0, 1); zeros(0, 1); zeros(0, 1); ...
            };
            for i = 1:length(keys)
                key = keys{i};
                default_value = default_values{i};
                if  ~isfield(data, key) || isempty(data.(key))
                    data.(key) = default_value;
                else
                    data.(key) = double(data.(key));
                end
            end
        end

        % %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% %
        % -=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=- %
        %%                          Dimension                          %%
        % -=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=-=x=- %
        % %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% %
        function n = dimension(key, data)
            % Key is the function/variable name
            % Value is it's dimension
            n = dictionary( ...
                'x', data.n_x, ...
                'y', data.n_y, ...
                'F', 1, ...
                'G', length(data.ineq_upper) + length(data.rhs_upper) + length(data.x_upper_bound) + length(data.x_lower_bound), ...
                'H', 0, ...
                'f', 1, ...
                'g', length(data.ineq_lower) + length(data.rhs_lower) + length(data.y_upper_bound) + length(data.y_lower_bound), ...
                'h', 0 ...
            );
            if isKey(n,key)
                n = n(key);
            end
        end
    end
end