classdef polynomial_bilevel
    % Polynomial Bilevel Program (PBP)
    % ================================
    % The objective function and constraints are polynomials.
    % Each polynomial is parameterised by m, c, p, q as defined below.
    % These parameters can be loaded from the JSON files in bolib3/data/polynomial_bilevel.
    %
    %
    % Polynomial expression
    % =====================
    %     m: Number of terms.
    %     c: Vector (m,) of coefficients.
    %     p: Matrix (m, n_x) of exponents for x.
    %     q: Matrix (m, n_y) of exponents for y.
    %
    %     In LaTeX syntax:
    %         Polynomial(x, y) := sum_{t=1,...,m}(
    %                                 c_t
    %                                 * product_{i=1,...,n_x}( x_i^{p_{t,i}} )
    %                                 * product_{i=1,...,n_y}( y_i^{q_{t,i}} )
    %                             )
    %
    %     In MATLAB syntax:
    %         Polynomial(x, y) :=  sum(c .* prod(x'.^p, 2) .* prod(y'.^q, 2))


    properties(Constant)
        name = 'polynomial_bilevel';
        category = 'archetype';
        subcategory = '';
        author = 'Samuel Ward';
        datasets = {
            'mitsos_barton_2006_ex310.json';
            'mitsos_barton_2006_ex311.json';
            'mitsos_barton_2006_ex312.json';
            'mitsos_barton_2006_ex313.json';
            'mitsos_barton_2006_ex314.json';
            'mitsos_barton_2006_ex315.json';
            'mitsos_barton_2006_ex316.json';
            'mitsos_barton_2006_ex317.json';
            'mitsos_barton_2006_ex318.json';
            'mitsos_barton_2006_ex319.json';
            'mitsos_barton_2006_ex320.json';
            'mitsos_barton_2006_ex326.json';
        };
        paths = fullfile('bolib3', 'data', 'polynomial_bilevel', polynomial_bilevel.datasets);
    end

    methods(Static)

        function val = polynomial(x, y, coeff_expon)
            % Poynomial in
            %       x: vector (n_x, 1)
            %       y: vector (n_x, 1)
            % The coeff_expon argument should be a stuct with fields
            %       c: vector (m, 1)   of coefficients.
            %       p: matrix (m, n_x) of exponents for x.
            %       q: matrix (m, n_y) of exponents for y.
            val = sum(coeff_expon.c .* ...
                prod(power(x(:).', coeff_expon.p), 2) .* ...
                prod(power(y(:).', coeff_expon.q), 2) ...
            );
        end

        function val = F(x, y, data)
            % Upper-level objective function
            val = polynomial_bilevel.polynomial(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) = polynomial_bilevel.polynomial(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 = polynomial_bilevel.polynomial(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) = polynomial_bilevel.polynomial( 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 = polynomial_bilevel.polynomial_derivative(x, y, data.F, 'x');
        end

        function dy = dF_dy(x, y, data)
            % ∂F/∂y
            % Dimention (n_y × 1)
            dy = polynomial_bilevel.polynomial_derivative(x, y, data.F, 'y');
        end

        function Hxx = d2F_dxx(x, y, data)
            % ∂²F/∂x²
            % Dimention (n_x × n_x)
            Hxx = polynomial_bilevel.polynomial_derivative(x, y, data.F, 'xx');
        end

        function Hyy = d2F_dyy(x, y, data)
            % ∂²F/∂y²
            % Dimention (n_y × n_y)
            Hyy = polynomial_bilevel.polynomial_derivative(x, y, data.F, 'yy');
        end

        function Hxy = d2F_dxy(x, y, data)
            % ∂²F/∂x∂y
            % Dimention (n_x × n_y)
            Hxy = polynomial_bilevel.polynomial_derivative(x, y, data.F, 'xy');
        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)
                    Jx(i,:) = polynomial_bilevel.polynomial_derivative( ...
                        x, y, data.ineq_upper(i), 'x' ...
                    );
                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)
                    Jy(i,:) = polynomial_bilevel.polynomial_derivative( ...
                        x, y, data.ineq_upper(i), 'y' ...
                    );
                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 = polynomial_bilevel.polynomial_derivative(x, y, data.f, 'x');
        end

        function dy = df_dy(x, y, data)
            % ∂f/∂y
            % Dimention (n_y × 1)
            dy = polynomial_bilevel.polynomial_derivative(x, y, data.f, 'y');
        end

        function Hxx = d2f_dxx(x, y, data)
            % ∂²f/∂x²
            % Dimention (n_x × n_x)
            Hxx = polynomial_bilevel.polynomial_derivative(x, y, data.f, 'xx');
        end

        function Hyy = d2f_dyy(x, y, data)
            % ∂²f/∂y²
            % Dimention (n_y × n_y)
            Hyy = polynomial_bilevel.polynomial_derivative(x, y, data.f, 'yy');
        end

        function Hxy = d2f_dxy(x, y, data)
            % ∂²f/∂x∂y
            % Dimention (n_x × n_y)
            Hxy = polynomial_bilevel.polynomial_derivative(x, y, data.f, 'xy');
        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)
                    Jx(i,:) = polynomial_bilevel.polynomial_derivative( ...
                        x, y, data.ineq_lower(i), 'x' ...
                    );
                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)
                    Jy(i,:) = polynomial_bilevel.polynomial_derivative( ...
                        x, y, data.ineq_upper(i), 'y' ...
                    );
                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


        function out = polynomial_derivative(x, y, coeff_expon, with_respect_to)
            %   Compute first/second derivatives of a mixed polynomial in x and y.
            %   
            %   Inputs
            %   ------
            %   x : vector (n_x, 1)
            %   y : vector (n_y, 1)
            %   coeff_expon : struct with fields
            %       .c : (m, 1) vector of coefficients
            %       .p : (m, n_x) exponents for x
            %       .q : (m, n_y) exponents for y
            %   with_respect_to : char, one of 'x', 'y', 'xx', 'xy', 'yy'
            %       'x'  -> gradient wrt x  (n_x, 1)
            %       'y'  -> gradient wrt y  (n_y, 1)
            %       'xx' -> Hessian wrt x   (n_x, n_x)
            %       'xy' -> cross Hessian   (n_x, n_y)
            %       'yy' -> Hessian wrt y   (n_y, n_y)
            %
            %   Output
            %   ------
            %   out : derivative (vector or matrix) as described above.

            arguments
                x double {mustBeVector}
                y double {mustBeVector}
                coeff_expon struct
                with_respect_to char {mustBeMember(with_respect_to, {'x','y','xx','xy','yy'})}
            end

            % Normalize shapes
            x_row = x(:).';   % (1, n_x)
            y_row = y(:).';   % (1, n_y)
        
            % Extract coefficients and exponents
            c = coeff_expon.c(:);    % (m x 1)
            p = coeff_expon.p;       % (m x n_x)
            q = coeff_expon.q;       % (m x n_y)
        
            [~, n_x] = size(p);  assert(n_x ~= 0);
            [~, n_y] = size(q);  assert(n_y ~= 0);


            % Precompute per-term base powers 
            % And then take their product along the second axis
            x_pow      = power(x_row, p);     % (m, n_x)
            y_pow      = power(y_row, q);     % (m, n_x)
            y_pow_prod = prod(y_pow, 2);  % (m, 1)
            x_pow_prod = prod(x_pow, 2);  % (m, 1)

            % Row-wise product excluding each column one by one.
            x_pow_exc_prod = polynomial_bilevel.exclusive_row_products(x_pow);
            y_pow_exc_prod = polynomial_bilevel.exclusive_row_products(y_pow);
        
            switch with_respect_to

                % Gradient with respect to x: (n_x, 1)
                case 'x'        
                    out = zeros(n_x, 1);
                    for k = 1:n_x
                        term = c .* y_pow_prod .* p(:,k) .* ...
                               x_pow_exc_prod(:,k) .* ( x_row(k) .^ max(p(:,k)-1, 0) );
                        out(k) = sum(term);
                    end
                
                % Gradient with respect to y: (n_y, 1)
                case 'y'
                    out = zeros(n_y, 1);
                    for j = 1:n_y
                        term = c .* x_pow_prod .* q(:,j) .* ...
                               y_pow_exc_prod(:,j) .* ( y_row(j) .^ max(q(:,j)-1, 0) );
                        out(j) = sum(term);
                    end
                
                % Hessian with respect to x twice: (n_x, n_x)
                case 'xx'
                    out = zeros(n_x, n_x);
                    for k = 1:n_x
                        for l = 1:n_x
                            if k == l
                                expk = p;
                                expk(:,k) = max(p(:,k) - 2, 0);
                                row_val = c .* y_pow_prod .* (p(:,k) .* max(p(:,k)-1, 0)) .* ...
                                          prod(power(x, expk), 2);
                            else
                                expkl = p;
                                expkl(:,k) = max(p(:,k) - 1, 0);
                                expkl(:,l) = max(p(:,l) - 1, 0);
                                row_val = c .* y_pow_prod .* (p(:,k) .* p(:,l)) .* ...
                                          prod(power(x, expkl), 2);
                            end
                            out(k,l) = sum(row_val);
                        end
                    end
                
                % Hessian with respect to y twice. (n_x, n_y)
                case 'yy'
                    out = zeros(n_y, n_y);
                    for j = 1:n_y
                        for h = 1:n_y
                            if j == h
                                expj = q;
                                expj(:,j) = max(q(:,j) - 2, 0);
                                row_val = c .* x_pow_prod .* (q(:,j) .* max(q(:,j)-1, 0)) .* ...
                                          prod(power(y, expj), 2);
                            else
                                expjh = q;
                                expjh(:,j) = max(q(:,j) - 1, 0);
                                expjh(:,h) = max(q(:,h) - 1, 0);
                                row_val = c .* x_pow_prod .* (q(:,j) .* q(:,h)) .* ...
                                          prod(power(y, expjh), 2);
                            end
                            out(j,h) = sum(row_val);
                        end
                    end
                   
                % Hessian with respect to x and y. (n_x, n_y)
                case 'xy'
                    out = zeros(n_x, n_y);
                    for k = 1:n_x
                        X_part = p(:,k) .* x_pow_exc_prod(:,k) .* ( x(k) .^ max(p(:,k)-1, 0) );
                        for j = 1:n_y
                            Y_part = q(:,j) .* y_pow_exc_prod(:,j) .* ( y(j) .^ max(q(:,j)-1, 0) );
                            row_val = c .* (X_part .* Y_part);
                            out(k,j) = sum(row_val);
                        end
                    end
                
                % Otherwise
                otherwise
                    error('Unsupported derivative type: %s', with_respect_to);
            end
        end


        function out = exclusive_row_products(matrix)
            % Row-wise product excluding each column one by one.
            %
            %   matrix: (m x n)
            %   out:    (m x n),   out(i,j) = prod_{k:k!=j} matrix(i,k)
            %
            [m, n] = size(matrix);
            L = ones(m, n);
            if n > 1
                L(:, 2:end) = cumprod(matrix(:, 1:end-1), 2);
            end
            R = ones(m, n);
            if n > 1
                R(:, 1:end-1) = fliplr(cumprod(fliplr(matrix(:, 2:end)), 2));
            end
            out = L .* R;
        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 and exponents to doubles
            keys = {'F', 'f'};
            subkeys = {'c', 'p', 'q'};
            for i = 1:length(keys)
                key = keys{i};
                % Diagnostics message
                msg = sprintf(strcat('Invalid polynomial for %s. ', ...
                    'It should look like this ', ...
                    '\n    "%s": { "c": [[]], "p": [[]], "q": [[]] }. ', ...
                    '\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_map = 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_map, key)
                n = n_map(key);
            else
                n = n_map;
            end
        end
    end
end