%EUSIPCO 2012
%Imperial College London 
%Andriy Gelman, Jon Onativia and Pier Luigi Dragotti 


function [optimum_pass] = lagrangian_multipliers_v2(rate, distortion, passes, desired_rate, index_distortion, total_number_of_blocks)
    
    %modify the rate into bits
    for i = 1:length(rate)
        rate{i} = rate{i} * 8;
    end
    
    %Choose two value for lamda, and obtain the optimum value by disection.  
    %Initialize
    lamda(1) = 1e50; lamda(3) = 0.001; lamda(2) = (lamda(1)+lamda(3))/2;
    total_rate(1) =0;
    total_rate(2) =0;
    total_rate(3) =0; %Total rate defines the optimum rate for the corresponding lamda

    %For each lamda evaluate the optimum rate
    optimum_rate = zeros(3, total_number_of_blocks);
    optimum_distortion = zeros(3, total_number_of_blocks);
    optimum_pass = zeros(3, total_number_of_blocks);


    while abs((lamda(1)-lamda(3))/lamda(1)) > 1e-10

        for i = 1:length(index_distortion)
            %i
            [optimum_rate(1,index_distortion(i)), optimum_distortion(1,index_distortion(i)), optimum_pass(1,index_distortion(i))] = choose_optimum_rate_v2(rate{index_distortion(i)}, distortion{index_distortion(i)}, passes{index_distortion(i)}, lamda(1));
            [optimum_rate(2,index_distortion(i)), optimum_distortion(2,index_distortion(i)), optimum_pass(2,index_distortion(i))] = choose_optimum_rate_v2(rate{index_distortion(i)}, distortion{index_distortion(i)}, passes{index_distortion(i)}, lamda(2));
            %if i==668
            %    test=1;
            %end
            %i
            [optimum_rate(3,index_distortion(i)), optimum_distortion(3,index_distortion(i)), optimum_pass(3,index_distortion(i))] = choose_optimum_rate_v2(rate{index_distortion(i)}, distortion{index_distortion(i)}, passes{index_distortion(i)}, lamda(3));
        end

    total_rate(1) = sum(optimum_rate(1,:));
    total_rate(2) = sum(optimum_rate(2,:));
    total_rate(3) = sum(optimum_rate(3,:));
    
    
    if total_rate(2) > desired_rate
        lamda(3) = lamda(2);
        lamda(2) = (lamda(1)+lamda(3))/2;
    else
        lamda(1) = lamda(2);
        lamda(2) = (lamda(1) + lamda(3))/2;
    end
    
    %if (abs(total_rate(1)-total_rate(3)) < 50)
    %    break;
    %end
    
    end
    
    %choose optimal lamda which has the lowest cost
    
    %apply second optimization stage which allocates left over bits which
    %will decrease the distortion the most
    
    optimum_pass = optimum_pass(1,:);
    optimum_rate = optimum_rate(1,:);
    total_rate = total_rate(1);
    
    %cost = abs(total_rate - desired_rate);
    %index = find(cost == min(cost), 1);
    
    while desired_rate > total_rate
        
        max_rate_change = desired_rate - total_rate;
        
        %for each RD evaluate the change in rate if pass is increamne
        change_in_rate = zeros(size(distortion));
        change_in_distortion = zeros(size(distortion));
        
        for l = 1:length(change_in_rate)
            index = find(passes{l} == optimum_pass(l));
            if isempty(index)
                change_in_rate(l) = 1e20;
            else
                change_in_rate(l) = rate{l}(index + 1) -  rate{l}(index);  
            end
        end
        
        check_rate_ind = find(change_in_rate < max_rate_change); 
        
        %change_in_rate(change_in_rate > max_rate_change) = 1e20;
        
        if isempty(check_rate_ind)
            %exit loop if no more bits can be allocated
            break;
        end
        
        
        
        for l = 1:length(change_in_distortion)
            index = find(passes{l} == optimum_pass(l));
            if isempty(index)
                change_in_distortion(l) = 0;
            else
               change_in_distortion(l) = distortion{l}(index) -  distortion{l}(index+1);  
            end

            
        end
        
        %find the point which reduces the smallest increase in rate and largest
        %increase in distortion
        deriv_grad_rate = change_in_distortion./change_in_rate;
        
        if max(deriv_grad_rate((check_rate_ind))) < 0
            break;
        end
        
        if isnan(max(deriv_grad_rate((check_rate_ind))))
            break;
        end
        
        
        ind_max_der = find(deriv_grad_rate(check_rate_ind) == max(deriv_grad_rate(check_rate_ind)), 1);
        buffer = deriv_grad_rate(check_rate_ind);
        
        ind_max_der = find(deriv_grad_rate == buffer(ind_max_der),1);
        
        index = find(passes{ind_max_der} == optimum_pass(ind_max_der));
        
        optimum_pass(ind_max_der) = optimum_pass(ind_max_der) + 1;
        optimum_rate(ind_max_der) = rate{ind_max_der}(index + 1);
        
        total_rate = sum(optimum_rate);
        
    end
    %output_lambda = lamda(index);
    %optimum_pass = optimum_pass(index,:);
end
