%% IPFP basic implementation %% JDB oct 2014
format long 
LP = 0 ; % do LP test 
IP = 1 ; % do IPFP test 

% Temperature relaxatin 
epsilon = 0.01 ;
% grid 
N = 1024 ; 
x = linspace(0,1,N) ;  % C 
y = linspace(0,1,N)'; 
%% tools 
oo = ones(1,N) ;

% cost and nargins 


c=@(x,y,oo) (oo'*x - y*oo).^2; %% quadratic 
%c=@(x,y,oo) sqrt((oo'*x - y*oo).^2); %% monge 
%c=@(x,y,oo) -(oo'*x).*(y*oo); %% quadratic 


marg1=@(x) exp(-10*(x-0.2).^2) ; 
marg2=@(y) exp(-10*(y-0.7).^2) ; 

%marg1=@(x) (x>0.1).*(x<0.3); 
%marg2=@(y)  (y>0.7).*(y<0.9)  ; 

%marg1=@(x) ones(size(x))  ; 
%marg2=@(y) ones(size(y)) ; 

% normalise
nu = marg1(x) ;
nu = N*nu/sum(nu) ; 
%
mu = marg2(y) ; 
mu = N*mu/sum(mu) ; 
%% plot margins 
figure(1) 
plot(x,nu,'b-', y,mu ,'r') ; 
legend('first margin', 'second margin') ; 

% Temperature relaxatin 

C = c(x,y,oo) ;

if (LP) %% do LP  

%%% linprog attempt  
f = C(:) ; %% cost 
beq = [nu(2:N) mu']' ; % vector of margins 
i = [1:N] ; j = [1:N^2] ; [I,J] = meshgrid(i,j) ; 
Aeq2 = (J>=(I-1)*N+1).*(J<= (I)*N)  ;
Aeq = [  Aeq2(:,2:N)' ; repmat(eye(N),1,N) ]; %% remove a line
options = optimset('Display','iter') ; 
tic 
[g,fval,exitf] = linprog(f,[],[],Aeq,beq,zeros(size(f)),[],[],options)  ;

'time for lp alg', toc 

Gamma_lp = reshape(g,N,N)  ;
figure(11) ; 
pcolor(x,y,Gamma_lp(:,:) ) ; 
shading flat ;
colorbar ; 
' LP result for distance  ' , fval/N^2
pause 

end 

if(IP)

%% KL Proj target 
H = exp(-C/epsilon) ; 

%% init IPFP 
b=ones(size(x)); %  L 
a=ones(size(y)) ;  %%  C 
an=a;

err=1;
count=1;
cmax = 1000 ; 

tic 

%% main loop 
while ( err>1e-03  &  count < cmax)  
    
%s = log(prod(exp(H.*(a*oo)),1))   ;
s = sum(H.*(a*oo),1)   ;
b = nu./s ;

s = sum(H.*(oo'*b),2) ;
an = mu./s;

err=sum(abs(an-a))/sum(a);

 str = sprintf('Error at step (%d): %s', count, err);
 disp(str);
 
 %% iterate
a=an;

count=count+1;
end

'time for ipfp alg ' , toc 

figure(2) 
Gamma = H.*(a*b) ;
pcolor(x,y,Gamma) 
shading flat 
colorbar

figure(3) 
plot(x,epsilon*log(b),'b-', y,epsilon*log(a) ,'r') ; 
legend('first potential', 'second potential') ; 

' IPFP result for distance  ' , Gamma(:)'*C(:)/N^2

end 
