%CPD_RIGID The rigid CPD point-set registration. It is recommended to use
%   and umbrella function rcpd_register with an option opt.method='rigid'
%   instead of direct use of the current funciton.
%
%
%   Input
%   ------------------
%   X, Y       real, double, full 2-D matrices of point-set locations. We want to
%              align Y onto X. [N,D]=size(X), where N number of points in X,
%              and D is the dimension of point-sets. Similarly [M,D]=size(Y).
%
%   rot=[0 or 1]    (default 1) 1- strict rotation, 0- allow reflections.
%   scale=[0 or 1]  (default 1) 1- estimate scaling, 0- don't estimate scaling.    
%   max_it          maximum number of iterations, try 150
%   tol             tolerance criterium, try 1e-5
%   viz=[0 or 1]    Visualize every iteration         
%   outliers=[0..1] The weight of noise and outliers, try 0.1
%   fgt=[0 or 1]    (default 0) Use a Fast Gauss transform (FGT). (use only for the large data problems)
%   corresp=[0 or 1](default 0) estimate the correspondence vector.
%
%
%
%
%   Output
%   ------------------
%   C      Correspondance vector, such that Y corresponds to X(C,:).
%   R      Rotation matrix.
%   t      Translation vector.
%   s      Scaling constant.
%   sigma2 Final sigma^2
%   iter   Final number or iterations
%   T      Registered Y point set
%
%
%   Examples
%   --------
%   It is recommended to use an umbrella function cpd_register with an
%   option opt.method='rigid' instead of direct use of the current
%   funciton.
%
%   See also CPD_REGISTER

% Copyright (C) 2008 Andriy Myronenko (myron@csee.ogi.edu)
%
%     This file is part of the Coherent Point Drift (CPD) package.
%
%     The source code is provided under the terms of the GNU General Public License as published by
%     the Free Software Foundation version 2 of the License.
%
%     CPD package is distributed in the hope that it will be useful,
%     but WITHOUT ANY WARRANTY; without even the implied warranty of
%     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%     GNU General Public License for more details.
%
%     You should have received a copy of the GNU General Public License
%     along with CPD package; if not, write to the Free Software
%     Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

function [C, R, t, s, sigma2, iter, T, max_score,min_score,last_score,ntol,all_scores,ntols,sigmas,all_scores_mov_to_targ]=cpd_rigid(X,Y, rot, scale, max_it, tol, viz, outliers, fgt, corresp, sigma2,opt,optreg)


% opt will contain opt.num_x_orig, opt.num_y_orig

[N, D]=size(X);[M, D]=size(Y);
if viz, figure; end;

% Initialization
if ~exist('sigma2','var') || isempty(sigma2) || (sigma2==0), 
    sigma2=(M*trace(X'*X)+N*trace(Y'*Y)-2*sum(X)*sum(Y)')/(M*N*D);
    
end
sigma2_init=sigma2;

T=Y; s=1; R=eye(D);
Y_init = Y;
X_init = X;

% Optimization
iter=0; ntol=tol+10; L=0;
counter = 0;

max_score = 0;
min_score = 1E15; 
new_score = 0;
new_score_fwd = 0;
all_scores = 0;
ntols = 0;
sigmas = 0;


toggler =0 ;
while (iter<max_it) && (ntol > tol) && (sigma2 > 10*eps) && (counter<3)
    
    if(iter==1)
        figure;
    end
    
    L_old=L;
    sigma2_orig = sigma2;
    
    
    % Check wheather we want to use the Fast Gauss Transform
    if ((sigma2>=1.7) && fgt == 1) || (fgt == 0)
    if (fgt==0)  % no FGT
        [P1,Pt1, PX, L]=cpd_P(X,T, sigma2 ,outliers); st='';
    else         % FGT
        [P1, Pt1, PX, L, sigma2, st]=cpd_Pfast(X, T, sigma2, outliers, sigma2_init, fgt);
    end
    if (~isempty(find(isnan(P1)>0,1)) || ~isempty(find(isnan(Pt1)>0,1)) || ~isempty(find(isnan(PX)>0,1)) || ~isempty(find(isnan(L)>0,1)) || ~isempty(find(isnan(sigma2)>0,1)) || ~isempty(find(isnan(st)>0,1)))
        [P1,Pt1, PX, L]=cpd_P(X,T, sigma2_orig ,outliers); st='';
        fprintf('ran into trouble')
        counter = counter + 1;
    end



    ntol=abs((L-L_old)/L);
    disp([' CPD Rigid ' st ' : dL= ' num2str(ntol) ', iter= ' num2str(iter) ' sigma2= ' num2str(sigma2) ' current_total_score = ' num2str(new_score) ' forwards score ' num2str(new_score_fwd)]);     
    
    Np=sum(Pt1);
    mu_x=X'*Pt1/Np;
    mu_y=Y'*P1/Np;
    
    % Solve for Rotation, scaling, translation and sigma^2
    A=PX'*Y-Np*(mu_x*mu_y'); % A= X'P'*Y-X'P'1*1'P'Y/Np; 

A_alt = A;

init_weight = .65;
first_sigma = 20;
second_sigma = 5;

C = (init_weight - 1)/(first_sigma - second_sigma);

if sigma2 > first_sigma
    final_weight = init_weight;
    fprintf('ok')
elseif (sigma2 <= first_sigma) && (sigma2 > second_sigma)
    final_weight = C*(sigma2 - second_sigma) + 1
else
    final_weight = 1;
end

final_weight

    fw = final_weight;

    A_alt(1,2) = fw*A(1,2); A_alt(2,1) = fw*A(2,1); A_alt(3,1) = fw*A(3,1); A_alt(1,3) = fw*A(1,3); A_alt(3,2) = fw*A(3,2); A_alt(2,3) = fw*A(2,3);
    
    A_alt

        
    [U,S,V]=svd(A); C=eye(D); Calt = eye(D);
    [Ualt,Salt,Valt] = svd(A_alt);
    
    if rot, C(end,end)=det(U*V'); Calt(end,end)=det(Ualt*Valt'); end % check if we need strictly rotation (no reflections)
%     R=U*C*V';
    R = Ualt*Calt*Valt';
    
    sigma2save=sigma2;
    if scale  % check if estimating scaling as well, otherwise s=1
        s=trace(S*C)/(sum(sum(Y.^2.*repmat(P1,1,D))) - Np*(mu_y'*mu_y));
        sigma2=abs(sum(sum(X.^2.*repmat(Pt1,1,D))) - Np*(mu_x'*mu_x) -s*trace(S*C))/(Np*D);
%         sigma2_sp = abs(sum(sum(X_sp.^2.*repmat(Pt1_sp,1,D))) - Np_sp*(mu_x_sp'*mu_x_sp) -s*trace(S_sp*C))/(Np_sp*D);
    else
        sigma2=abs((sum(sum(X.^2.*repmat(Pt1,1,D))) - Np*(mu_x'*mu_x)+sum(sum(Y.^2.*repmat(P1,1,D))) - Np*(mu_y'*mu_y) -2*trace(S*C))/(Np*D));
        sigma2=abs((sum(sum(X.^2.*repmat(Pt1,1,D))) - Np*(mu_x'*mu_x)+sum(sum(Y.^2.*repmat(P1,1,D))) - Np*(mu_y'*mu_y) -2*trace(Salt*Calt))/(Np*D));
%         sigma2_sp = abs((sum(sum(X_sp.^2.*repmat(Pt1_sp,1,D))) - Np_sp*(mu_x'*mu_x)+sum(sum(Y_sp.^2.*repmat(P1_sp,1,D))) - Np_sp*(mu_y_sp'*mu_y_sp) -2*trace(S_sp*C))/(Np_sp*D));
    end
    
    t=mu_x-s*R*mu_y;
    
    % Update the GMM centroids
    T=s*Y*R'+repmat(t',[M 1]);
        
    iter = iter+1;

    if viz, cpd_plot_iter(X, T); end; % show current iteration if viz=1
    else
        fprintf('sigma 2 below min for fast gauss transform, quitting registration')
        ntol = tol-1;
    end
    
if corresp, C=cpd_Pcorrespondence(X,T,sigma2save,outliers); else C=0; end;

length_X = length(X);
length_T = length(T);
length_C = length(C);

% current_total_score = 0;
% for i= 1:5000
%     if C(i) <= 5000
%         current_total_score = current_total_score + norm(X(C(i),:) - T(i,:));
%     else
%         current_total_score = current_total_score + 100;
% % do nothing
%     end
% end
% 
% 

% new_score = 0;
% % new_scores = 0;
% for i = 1:5000
%     for j = 1:5000
%         vect = X(i,:)-T(j,:);
%         new_score = exp((-sum((X(i,:)-T(j,:)).^2))/36) + new_score;
%     end
% end

size(X)
size(T)
new_score= TestScoreCalculator(X(1:5000,:),T(1:5000,:));
new_score_fwd = TestScoreCalculator(T(1:5000,:),X(1:5000,:));
all_scores(iter) = new_score;
all_scores_mov_to_targ(iter) = new_score_fwd;
ntols(iter) = ntol;
sigmas(iter) = sigma2;

% if toggler>0
    max_score = max(new_score,max_score);
    min_score = min(new_score,min_score);
% else
    toggler = 1;
    
% end

end
last_score = new_score;