% function [S1, S2, b1, b2] = optimize_s(Hmag_sq, No, w, ...
%          lambda_1, lambda_2, mask, gap, bitcap, cb)
%
% OSB tonal Lagrangian tests values 0:bitcap for users 1 and 2. 
% A. Chowdhery ~2010 ; Updated by J. Cioffi in 2024. This is the OSB 
% computational subroutine. Again, original by A. Chowdhery
% and ppdated by J. Cioffi in 2024.  This program presently handles
% only 2 users, so U=2. 
%
% Inputs 
%
% Hmag_sq is a N x 2 x 2 where N is FFT size. N inferred from this.
% No      is a 1 x U white-noise power spectra density matrix.
%    If Hmag_sq is complex BB, then No should be the one-sided PSD.
% E2      is user 2's energy.
% theta   is a 1 x U user-rate weighting vector.
% w_1      is user 1's energy weighting factor.
% w_2      is user 2's energy weighting factor.
% mask    is an N x U PSD maximum allowed.
% gap     is the (non-dB) linear gap (so 1 if 0 dB gap).
% bitcap  is a 1 x U maximum number of bits allowed per tone.
% cb      is 2 for real baseband and 1 for cplex bband
%
% Outputs
%
% S1      is user 1's Nx1 PSD
% S2      is user 2's Nx1 PSD
% b1      is user 1's Nx1 bit distribution
% b2      is user 2's Nx1 bit distribution
%
% -------------------------------------------------------------
function [S1, S2, b1, b2] = optimize_s(Hmag_sq, No, theta, w_1, w_2, mask, gap, bitcap, cb)

N = size(Hmag_sq,1);
S1 = zeros(N,1); S2 = zeros(N,1);
b1 = zeros(N,1); b2 = zeros(N,1);

for n = 1:N
    max_value = -inf;
    for i1 = 0:bitcap(1)
        for i2 = 0:bitcap(2)
            % next 3 lines essentially invert a matrix to find the psds
            D = Hmag_sq(n,1,1)*Hmag_sq(n,2,2)-gap^2*(2^i1-1)*(2^i2-1)*Hmag_sq(n,1,2)*Hmag_sq(n,2,1);
            N1 = gap*(2^i1-1)*(gap*(2^i2-1)*Hmag_sq(n,2,1)*No(n,2)+Hmag_sq(n,2,2)*No(n,1));
            N2 = gap*(2^i2-1)*(gap*(2^i1-1)*Hmag_sq(n,1,2)*No(n,1)+Hmag_sq(n,1,1)*No(n,2));
            s(1) = N1/D; s(2) = N2/D;
            Lk = theta*i1 + (1-theta)*i2 - w_1*s(1) - w_2*s(2);
            if Lk > max_value & s(1)<=mask(n,1) & s(2)<=mask(n,2) & s(1)>=0 & s(2)>=0
                max_value = Lk;
                b1(n) = i1/cb; b2(n) = i2/cb;
                S1(n) = s(1); S2(n) = s(2);
            end
        end
    end
end

%fprintf('w = %3.2f, l1 = %5.1e, l2 = %5.1e, b1 = %4d, b2 = %4d, E1 = %5.4f, E2 = %5.4f \n',...
 %   theta, lambda_1, lambda_2, sum(b1), sum(b2), sum(S1), sum(S2) );
%drawnow;

return;