function [reg_p nrg] = imreg (i1, i2, K, method, lb, ub)

%PUBLISHED UNDER THE CREATIVE COMMONS LICENSE (SEE MY WEBPAGE FOR MORE DETAILS)

%rigid image registration function
%INPUT:
% i1,i2 are the two images to register
% K is the number of clusters in both images
% lb,ub are the lower, upper bounds respectively for the parameters space
%  lb, ub = [theta,tx,ty];
%  with X0 = (ub+lb)/2 
%
%  in case ub is empty, unconstraint optimization considered, with X0 = lb
% method: which method to follow for registration:
%  MI -> mutual information
%  GMM -> registration with Gaussian Mixture Models (GMM)
%  tMM -> registration with Student's t Mixture Models (TMM)
%  CORR -> registration using correlation of the pixels
%
%
%OUTPUT:
%reg_p = [theta,tx,ty] -> registration parameters estimated (theta in degrees, tx,ty in pixels)
%reg_p = registration energy

%OPTIMIZATION PROPERTIES
optim_par1 = optimset ('TolX',1E-02,'TolFun',1E-04,...
    'MaxFunEvals',1000, 'Display','Off');

optim_par2 = optimset ('TolX',1E-02,'TolFun',1E-04,...
    'Algorithm','active-set',...
    'LargeScale','Off','GradObj','Off',...
    'MaxFunEvals',1000, 'Display','Off');

OPTIM_ITER = 1;

if isempty (lb)
    error ('Lower bound, or initial point is not specified.');
end

%GET THE OBJECTIVE FUNCTION (ENERGY)
if strcmp (method, 'MI')
    fun = @(w)imreg_nrg('MI',w,i2,i1,K);
elseif strcmp (method, 'GMM')
    %PIXELS WITH BLACK COLOR CONSIDERED AS CROPED AREAS - DROP THEM
    X = i2(:); ind = find (X~=0); mask = zeros (size(X));
    [mask1 p m s] = gmm_classifier (X(ind),K);
    mask(ind) = mask1;
    
    fun = @(w)imreg_nrg('MM',w,mask,p,m,s,i1);   
elseif strcmp (method, 'tMM')
    %PIXELS WITH BLACK COLOR CONSIDERED AS CROPED AREAS - DROP THEM
    X = i2(:); ind = find (X~=0); mask = zeros (size(X));
    [mask1 p m s v] = tmm_classifier (X(ind),K);
    mask(ind) = mask1;

    fun = @(w)imreg_nrg('MM',w,mask,p,m,s,i1);   
elseif strcmp(method, 'CORR')
    fun = @(w)imreg_nrg('CORR',w,i2,i1);
else
    error (sprintf ('%s: Uknown method.',method));
end


if strcmp(method, 'MI') || strcmp(method, 'CORR')
    % % DERIVATIVE FREE OPTIMIZATION (SIMPLEX)
    stp = mean(ub-lb)/OPTIM_ITER; nrg = inf;
    for i=1:OPTIM_ITER
        x0 = lb + (i-1)*stp;
        [reg_p1 nrg1] = fminsearch (fun,x0,optim_par1);
        if nrg1 < nrg
            nrg = nrg1;
            reg_p = reg_p1;
        end
    end
else
    %QUASI NEWTON
    if isempty(ub)
        %UNCONSTRAINT
        [reg_p nrg] = fminunc (fun,lb,optim_par2);
    else
        %CONSTRAINT
        stp = mean(ub-lb)/OPTIM_ITER; nrg = inf;
        for i=1:OPTIM_ITER
            x0 = lb + (i-1)*stp;
            [reg_p1 nrg1] = fmincon (fun,x0,[],[],[],[],lb,ub,[],optim_par2);
            if nrg1 < nrg
                nrg = nrg1;
                reg_p = reg_p1;
            end
        end
    end
end
%-

%-
% imreg_nrg
%-
function [nrg] = imreg_nrg (method,p,arg1,arg2,arg3,arg4,arg5)

% -----------------------------------------------------------------------
% CALCULATE THE ENERGY FUNCTION FOR A SPECIFIC IMAGE REGISTRATION
% TECHNIQUE.
%
% imreg_nrg (method,x1,x2,x2,x4)
%  method: Specifies the method for wich the energy is calculated.
%          Supported methods: MI,GMM-BATT,GMM-L2,tMM-BATT,tMM-L2
%       p: The configuration uppon which the energy is calculated.
%          (Rigid transformation parameters => Rotation, Translation
%          ! p = [Theta; Tx; Ty]. Theta in degrees, Tx,Ty in pixels !
%  The rest arguments are used according to the specified method. See
%  the code for more details.
% -----------------------------------------------------------------------

if strcmp (method,'MI')  
    % arg1:i_ref, arg2:i_reg, arg3:b
    % The rest arguments are ignored !
   
    %OBJECTIVE FUNCTION
    %perform the transformation for the input parameters p
    arg2 = imtrans (arg2,p);
      
    [h1 h2 hj] = get_histograms (arg1,arg2,arg3);
    
    T = sum(h1); h1 = h1 / (T + (T==0));
    H1 = - sum (h1.*log2(h1 + (h1==0)*eps)); %MARGINAL ENTROPY OF Iref

    T = sum(h2); h2 = h2 / (T + (T==0));
    H2 = - sum (h2.*log2(h2 + (h2==0)*eps)); %MARGINAL ENTROPY OF Ireg
    
    T = sum(hj(:)); hj = hj / (T + (T==0));
    H3 = - sum (sum((hj.*log2(hj + (hj==0)*eps)))); %JOINT ENTROPY

    %negative normalized mi (GOAL: maximize Mutual Information =>
    % minimize negative Mutual Information)
    nrg = -(H1+H2) / (H3 + (H3==0)*eps);
elseif strcmp (method,'MM')
    %arg1:mask_ref, arg2:p_ref, arg3:m_ref, arg4:s_ref, arg5:i_ref
   
    %perform the transformation for the input parameters p
    arg55 = imtrans (arg5,p); 
   
    arg55 = arg55(:);
    
    K = length(arg2); nrg = 0;
    for i=1:K       
        T = arg55.*(arg1==i); 
        T = T(find(T~=0));
        if isempty (T) continue; end
        
        g = std (T);
        if g < 1. %If cluster is too compact perhaps segmenation is not good
            % => penalize it with huge std.
            g = 100.;
        end
        %Modified Battacharrya distance - see paper(s) for details
        nrg = nrg + arg2(i) * (arg4(i)^2+g^2) / (arg4(i)*g);
    end
elseif  strcmp (method,'CORR')
    %arg1:i_ref, arg2:i_reg
   
    %perform the transformation for the input parameters p
    arg2 = imtrans (arg2,p);

    %get overlap
    arg1 = arg1.*(arg2 ~= 0);
    arg2 = arg2.*(arg1 ~= 0);
    
    %substruct overlap mean value
    ixx = arg1(:); ind = find(ixx); m = mean (ixx(ind));
    ixx(ind) = ixx(ind)-m; arg1 = reshape(ixx,size(arg1));
    
    ixx = arg2(:); ind = find(ixx); m = mean (ixx(ind));
    ixx(ind) = ixx(ind)-m; arg2 = reshape(ixx,size(arg2));
    
    
    arg11 = arg1.*arg2; arg11 = arg11(:); arg11 = arg11(find(arg1));
    
    arg1 = arg1(:); arg1 = arg1(find(arg1));
    arg2 = arg2(:); arg2 = arg2(find(arg2));
    
    %correlation - the goal is to maximize correlation,
    %but we use minimization algorithms
    nrg = -mean(arg11)/(std(arg1)*std(arg2));
end

%-
% get_histograms
%-
function [h1 h2 hj] = get_histograms (i1,i2, b)

% -----------------------------------------------------------------------
% get_histograms: get the marginal and joint histograms of i1,i2 images.
%                 ! HISTOGRAMS ARE CALCULATED ON OVERLAP !
%                 ! PIXELS WITH 0 VALUE ARE IGNORED !
%                 ! ALL IMAGE VALUES ARE POSITIVE (OR ZERO) !
%               !!! THE SIZE OF BOTH IMAGES (i1,i2) IS THE SAME !!!
% -----------------------------------------------------------------------
 %IMPORTANT: IMAGES HAVE THE SAME SIZE !!!
% -----------------------------------------------------------------------

%VECTORIZATION
i1 = i1(:); i2 = i2(:);

%FIND OVERLAP (=> MAKE PROJECTION)
%NON OVERLAPED AREAS ARE MARKED WITH ZERO
ind = find (i1 == 0); i2(ind) = 0;
ind = find (i2 == 0); i1(ind) = 0;

i1 = i1(find(i1~=0));
i2 = i2(find(i2~=0));

%Joint histogram
hj = zeros (b);
ma=min(i1);
MA=max(i1);
mb=min(i2);
MB=max(i2);

i1=round((i1-ma)*(b-1)/(MA-ma));
i2=round((i2-mb)*(b-1)/(MB-mb));

for i=0:b-1
    hj(i+1,:)=hist(i2(find(i1==i)),0:b-1);
end

%Marginal histograms 
% => marginalization of joint
h1=sum(hj,2);
h2=sum(hj,1)';