
Projs = importdata('DENOISED_smallN1_30mrad_DataMA_BM3D100.mat');
Projs2T = zeros(size(Projs,1)*2,size(Projs,2)*2,size(Projs,3));
for i=1:size(Projs,3)
    FFD = Projs(:,:,i);
    FFD2 = XTinterp(FFD,2);
    Projs2T(:,:,i)=FFD2;
end

Projs = Projs2T;

%% choose a region as reference
cropVectors = zeros(size(Projs,3),2);

figure;hold on;colormap gray
for i=1 :size(Projs,3)
    II = (squeeze(Projs(:,:,i)));

    while 1
        coordinates = [];
        My_clickpos(II)
        pause
        if isempty(coordinates) % if any key pressed without selecting some position
            prompt = 'select center!\n ';
        else
            break
        end
    end
    a1 = round(coordinates(2))            
    a2 = round(coordinates(1))
    cropVectors(i,1) = a1;
    cropVectors(i,2) = a2;

end
%% check the region
% load('cropVectors.mat')

figure;colormap('gray')
for i=1:size(Projs,3)
    imagesc(Projs(:,:,i));axis image;hold on;
    plot(cropVectors(i,2),cropVectors(i,1),'wo')
    hold off
    pause
end

%%
load('Model_ORNL_100_forDeconv.mat');
Angles=importdata('angles2.mat'); %in Fourier space, applied to coordinates
res = 0.3433;
f_pos = f_pos./res; % convert into pixels

comX = sum(f_pos(1,:))/numel(f_pos(1,:));
comY = sum(f_pos(2,:))/numel(f_pos(2,:));
f_pos(1,:) = f_pos(1,:)-comX;
f_pos(2,:) = f_pos(2,:)-comY;

Z=[0;0;1];


ROTZ=18.4349-1.5953-1.9160;
inprotM = MatrixQuaternionRot(Z,ROTZ); %MatrixQuaternionRot assumes ndgrid as coordinates

f_posA = inprotM*f_pos; %real space

proj=12;
Rotangles = Angles(proj,:);
rotM = MatrixQuaternionRot([0 0 1],Rotangles(1))*MatrixQuaternionRot([0 1 0],Rotangles(2))*MatrixQuaternionRot([1 0 0],Rotangles(3));

f_posFinal = rotM'*f_posA; %here the ZYX Euler angles are for coordinates rotation. For real space, use transpose

min(f_posFinal(1,:))
max(f_posFinal(1,:))
min(f_posFinal(2,:))
max(f_posFinal(2,:))
%%
f_posFinalAngstrom = f_posFinal.*res;
Widthar = ([0.552 0.552 0.53]).^2;
Magar = [ 0.4502    0.1369 73];  %doping with Rhenium

x0 = -59.5:0.5:60;
% z0 = -30:1:30;
[X,Y] = ndgrid(x0,x0);
sizeX = size(X);
X = X(:);
Y = Y(:);
% Z = Z(:);

finalvol = zeros(1,length(X));
for i=1:length(X)

%     finalvol(i) = sum( Magar(f_att).*exp(-1*( (X(i)*res-f_posFinalAngstrom(1,:)).^2 + (Y(i)*res-f_posFinalAngstrom(2,:)).^2 + (Z(i)*res-f_posFinalAngstrom(3,:)).^2)./(2*Widthar(f_att))),2);
    finalvol(i) = sum( Magar(f_att).*exp(-1*( (X(i)*res-f_posFinalAngstrom(1,:)).^2 + (Y(i)*res-f_posFinalAngstrom(2,:)).^2)./(2*Widthar(f_att))),2);

    if mod(i,100000) == 0
        i
    end
end

finalvol = reshape(finalvol,sizeX);

sumZ=sum(finalvol,3);
figure;imagesc(sumZ);axis image

%%
arr=-119:120;
imageRef=Projs(cropVectors(proj,1)+arr,cropVectors(proj,2)+arr,proj);
imageRef= imageRef - min(imageRef(:)) ;
centerX = 121;
centerY = 121;

windowhalfsizeX = 100;
windowhalfsizeY = 100;
offCenter = 0;

windowDim1ind = centerX+offCenter-windowhalfsizeX:centerX+offCenter+windowhalfsizeX;
windowDim2ind = centerY+offCenter-windowhalfsizeY:centerY+offCenter+windowhalfsizeY;

Res = 0.05;
searchRange1 = 10;
searchRange2 = 10;

[d1, d2, maxVal] = My_NormXcorr_subpixel_align_setWindow_atomicRegion(imageRef, sumZ, windowDim1ind, windowDim2ind, Res,searchRange1,searchRange2);
f_posFinal = rotM'*f_posA;
% f_posFinal(1,:)=f_posFinal(1,:)+d1 +centerX-1;
% f_posFinal(2,:)=f_posFinal(2,:)+d2 + centerY-1;
f_posFinal(1,:)=f_posFinal(1,:)+d1/2;
f_posFinal(2,:)=f_posFinal(2,:)+d2/2;

newsumZ = real(My_FourierShift(sumZ,d1,d2));

newsumZ = newsumZ - min(newsumZ(:)) ;
newsumZ = newsumZ / max(newsumZ(:)) ;
imageRefn = imageRef - min(imageRef(:)) ; % subtract background to get more reasonable model
imageRefn = imageRefn / max(imageRefn(:)) ; % normalization

loopShowTemp(imageRef+newsumZ)

%% masking zero proj
im=double(newsumZ/max(newsumZ(:))*255);
[ot,outim]=otsu_thresh(im);

ot = ot*max(max(newsumZ))/255*0.8;

% fill the gap
after_conv1 = (newsumZ>ot) * 1;       
after_conv1  = bwareaopen(after_conv1, 30); % some sparcles    
after_conv1 = imfill(after_conv1,'holes');
after_conv2 = bg_conv_gauss_corr(My_paddzero(after_conv1,size(after_conv1)+128), 2, 0.1);

se = strel('disk',4);
after_conv2 = My_stripzero(imerode(after_conv2,se,'same'),size(after_conv1));

imageRefnm=imageRefn.*after_conv2;
figure;subplot(121);imagesc(imageRefnm);axis image
subplot(122);imagesc(newsumZ);axis image
%% fit zero proj
imageRef=Projs(cropVectors(proj,1)+arr,cropVectors(proj,2)+arr,proj);
ydata = (imageRef - min(imageRef(:))).*after_conv2;

% ydata = My_stripzero(ydata,[120 120]);

f_posFinalAngstromFit = f_posFinal.*res;
xdata = [f_posFinalAngstromFit;f_att];

k0 = [0.5279    0.3434    0.4502    0.1369]; %init guess
lb = [0.3 0.3 0.01 0.01];
ub = [1 0.4 inf inf];

opt=optimset('TolFun',1e-7);
opt=optimset(opt,'Display','off');
[kFit, Resnorm ]= lsqcurvefit(@fit_toBuildModel_res_up2,k0,xdata,My_stripzero(ydata,[100 100]),lb,ub,opt)

figure;subplot(121);imagesc(My_stripzero(ydata,[100 100]));axis image
subplot(122);imagesc(My_stripzero(fit_toBuildModel_up2(kFit,xdata),[100 100]));axis image
[d1, d2, maxVal] = My_NormXcorr_subpixel_align_setWindow_atomicRegion(ydata, fit_toBuildModel_up2(kFit,xdata), windowDim1ind, windowDim2ind, Res,searchRange1,searchRange2);
[d1,d2] %to check whether any shift

%% kFit: 0.5859    0.6000    0.2163    0.0541 (BM3D300 Up2)

% kFit= [0.5291    0.3433    0.4881    0.1601]
% kFit= [0.5279    0.3434    0.4502    0.1369] %BM3D 300
% kFit= [0.5164    0.3434    0.4654    0.1417] %BM3D 200
% kFit= [0.5180    0.3435    0.4728    0.1449] %BM3D 150

kFit00 = kFit;

load('Model_ORNL_100_forDeconv.mat');
res = kFit00(2);
f_pos = f_pos./res; % convert into pixels


comX = sum(f_pos(1,:))/numel(f_pos(1,:));
comY = sum(f_pos(2,:))/numel(f_pos(2,:));
f_pos(1,:) = f_pos(1,:)-comX;
f_pos(2,:) = f_pos(2,:)-comY;

Z=[0;0;1];

inprotM = MatrixQuaternionRot(Z,ROTZ); %MatrixQuaternionRot assumes ndgrid as coordinates
f_posA = inprotM*f_pos; %real space

DeconvFit = Projs;
DeconvFit2 = Projs;

CORR2AR = zeros(size(Projs,3),50);
CORR2ARFinal = zeros(size(Projs,3),50);

kFinal = zeros(size(Projs,3),7);
kScanBest = zeros(size(Projs,3),3);

k2 = zeros(size(Projs,3),4); %for second time deconvolution
CORR2AR2 = zeros(size(Projs,3),50);

spatialFit = zeros(size(Projs,3),length(kFit00));
atomFit = zeros(size(Projs,3),length(kFit00));

simuStack = zeros(240,240,size(Projs,3));
simuStack2 = zeros(240,240,size(Projs,3));
CORR2Size = [100 100];

% kLast=load('Projs_Denoised_Deconv2_test14_HarmonicGaussiank7_ne.mat', 'kFinal')
% KKK = kLast.kFinal;
% kLast=load('ProjsNew1_DataMC_HarmGauss_k7_newKer_BM3D100.mat', 'kScanBest')
% KKK = kLast.kScanBest;

upSample = 2;
%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
normVectors=plotAnglePlane(Angles);
% note: normVectors is in NDGRID, but the kernels takes in meshgrid angles,
% so transform it to meshgrid


%%



for   proj = 1:size(Angles,1)

    imageRef=Projs(cropVectors(proj,1)+arr,cropVectors(proj,2)+arr,proj);
%     imageRef = imageRef - min(imageRef(:)); %background subtraction
    
    %align f_pos the data
    Rotangles = Angles(proj,:);
    rotM = MatrixQuaternionRot([0 0 1],Rotangles(1))*MatrixQuaternionRot([0 1 0],Rotangles(2))*MatrixQuaternionRot([1 0 0],Rotangles(3));
    f_posFinal = rotM'*f_posA; %here the ZYX Euler angles are for coordinates rotation. For real space, use transpose
    f_posFinalAngstrom = f_posFinal.*res;
    xdata1 = [f_posFinalAngstrom;f_att];
    simu1 = fit_toBuildModel_up2(kFit00,xdata1);
    [d1, d2, maxVal] = My_NormXcorr_subpixel_align_setWindow_atomicRegion(imageRef, simu1, windowDim1ind, windowDim2ind, Res,searchRange1,searchRange2);
    [d1,d2]
    f_posFinal(1,:)=f_posFinal(1,:)+d1/2;
    f_posFinal(2,:)=f_posFinal(2,:)+d2/2;
%     loopShow(imageRef+abs(My_FourierShift(simu1,d1,d2)));
    
    
    %correct rotation
    f_posFinalAngstrom = f_posFinal.*res;
    xdataNew = [f_posFinalAngstrom;f_att];
    simuNew = fit_toBuildModel_up2(kFit00,xdataNew);
    % fit to find the best inplane rotation
    [r,Resnom]=fit_findRotAndShift([0.5 0 0 1], simuNew, imageRef,[-5,-5,-5,0],[5,5,5,inf]); %k1 disabled
%     spatialFit(proj,:)=r;
    r
    Z=[0;0;1];
    inprotAdjust = MatrixQuaternionRot(Z,r(1)); %positive means anticlockwise
    f_posFinal = inprotAdjust*f_posFinal; %real space
    
    %align XY again
    f_posFinalAngstrom = f_posFinal.*res;
    xdataNew = [f_posFinalAngstrom;f_att];
    simuNew = fit_toBuildModel_up2(kFit00,xdataNew);
    [d1, d2, maxVal] = My_NormXcorr_subpixel_align_setWindow_atomicRegion(imageRef, simuNew, windowDim1ind, windowDim2ind, Res,searchRange1,searchRange2);
    [d1,d2]
    f_posFinal(1,:)=f_posFinal(1,:)+d1/2;
    f_posFinal(2,:)=f_posFinal(2,:)+d2/2;
    
    
    %masking using the updated model
    f_posFinalAngstrom = f_posFinal.*res;
    xdataNew = [f_posFinalAngstrom;f_att];
    newsumZ = fit_toBuildModel_up2(kFit00,xdataNew);
    newsumZ = newsumZ - min(newsumZ(:)) ;
    newsumZ = newsumZ / max(newsumZ(:)) ;
    im=double(newsumZ/max(newsumZ(:))*255);
    [ot,outim]=otsu_thresh(im);
    ot = ot*max(max(newsumZ))/255*0.6;

    after_conv1 = (newsumZ>ot) * 1;       
    after_conv1  = bwareaopen(after_conv1, 30); % some sparcles    
    after_conv1 = imfill(after_conv1,'holes');
    after_conv2 = bg_conv_gauss_corr(My_paddzero(after_conv1,size(after_conv1)+128), 2, 0.1);
    se = strel('disk',3);
    after_conv2 = My_stripzero(imerode(after_conv2,se,'same'),size(after_conv1));
    
    se = strel('disk',3);
    after_conv_small = My_stripzero(imerode(after_conv2,se,'same'),size(after_conv1));
    after_conv_small = conv2(after_conv_small,fspecial('gaussian',5,5),'same'); 
%     figure(11);subplot(121);imagesc(imageRef.*after_conv2);axis image;title('Mask2')
%     subplot(122);imagesc(imageRef.*after_conv_small);axis image;title('MaskSmall')
    
    imageRefFinal=imageRef.*after_conv_small;
%     figure(12);subplot(121);imagesc(imageRefFinal);axis image;title('Data')
%     subplot(122);imagesc(newsumZ);axis image;title('Model')
    
    % build ground truth atomic model
    f_posFinalAngstromShift = f_posFinal.*res;
    xdataFinal = [f_posFinalAngstromShift;f_att];
    simuFinal = fit_toBuildModel_up2(kFit00,xdataFinal);   
    
    xdata = cat(3,simuFinal,after_conv_small);
    
   %%
    bgMax = min(imageRef(:));
    
    fullAmplitude = 13;
    vec = normVectors(:,proj);
    angle1 = -atand(vec(2)/vec(1));
    amp = fullAmplitude*sind(acosd(vec(3)));

    % amplitude, angle,background,gaussianPerpendicular,
    % shiftX,shiftY,scaleFactor
    % twofoldAstigmationAngle,std1,std2
    
    initGuess = [ceil(amp),angle1,bgMax/2,0.1,0,0,1];
    
    if proj==1||proj==2||proj==4||proj==9||proj==10||proj==11
        lb = [1,-2,0,0.01,-2,-2,0.5];
        ub = [11,90,bgMax,10,2,2,1.5];
%         initGuess = [5,45,bgMax/2,0.05,0,0,1];
    elseif proj==3||proj==5||proj==6||proj==8||proj==7||proj==12
        lb = [1,-90,0,0.01,-2,-2,0.5];
        ub = [11,2,bgMax,10,2,2,1.5];
%         initGuess = [5,-45,bgMax/2,0.05,0,0,1];
    elseif proj==13
        lb = [1,-0,0,0.01,-2,-2,0.5];
        ub = [11,90,bgMax,10,2,2,1.5];
%         initGuess = [5,60,bgMax/2,0.05,0,0,1];
    end

%     initGuess(7)=[];
    %% find kernel by convolution
    % k: amplitude, angle, background, sigma, k(5:6)=shiftX shiftY,
    [k, kernel, finalFit] = fit_harmonicGaussian2D_Kernel_k7_new_up2(initGuess, xdata, imageRefFinal, [0  0 0 0 0 0 0],lb,ub); %k1 disabled

    k
    kFinal(proj,:)=k;
    
    simuShift = simuFinal; %abs(My_FourierShift(simuFinal,k(5) ,k(6)));
    simuStack(:,:,proj)=simuShift;
    
    figure(33);
    subplot(1,2,1)
    imagesc(kernel);axis image;title('kernel1')
    subplot(1,2,2)
    imagesc(My_stripzero(finalFit,size(simuShift)));axis image;title('simu')

    proj1 = (Projs(cropVectors(proj,1)+arr,cropVectors(proj,2)+arr,proj)-k(3))/k(7);
    WT = zeros(size(proj1));
    edge = 7;
    WT(edge:end-edge+1,edge:end-edge+1) = 1;
    
    deconvImage=proj1.*after_conv_small;
    CORR2AR(proj,1) = corr2(My_stripzero( deconvImage, CORR2Size),My_stripzero( simuShift,CORR2Size));
    
    pj1deconv  = deconvlucy({proj1},kernel,1,0,WT);
    deconvImage=pj1deconv{2}.*after_conv_small;
    CORR2AR(proj,2) = corr2(My_stripzero( deconvImage, CORR2Size),My_stripzero( simuShift,CORR2Size));
    
    pause(1)
    for iter = 3:50
        pj1deconv  = deconvlucy(pj1deconv,kernel,1,0,WT);
        deconvImage=pj1deconv{2}.*after_conv_small;
        CORR2AR(proj,iter) = corr2(My_stripzero( deconvImage,CORR2Size),My_stripzero( simuShift,CORR2Size));
    end
    
    figure(41);plot(CORR2AR(proj,:),'r-o');title(num2str(proj));pause(1)
    
    [maxValue, bestIteration]=max(CORR2AR(proj,:));
   
    if bestIteration==1 &&maxValue>0.96 
        pj1deconv = (Projs(:,:,proj)-k(3))/k(7);
%         DeconvFit(:,:,proj) = pj1deconv;
%         DeconvFit2(:,:,proj) = pj1deconv;
        loopShowTemp(pj1deconv);
        
%         continue
    elseif bestIteration==1||bestIteration>15
        bestIteration = 5;
    end
    kFinal
%     bestIteration
    bestIteration=5
        
    %% do looping to polish kernel (deconvolution step)
    fprintf('LOOPING FOR KERNEL\n');
    Amplitude = [kFinal(proj,1):2:kFinal(proj,1)+2];
    AngleInp = [kFinal(proj,2)-50:10:kFinal(proj,2)+50];
    SigmaKernel = [0.01:0.05:1.2];
    errorMatrix = zeros([length(Amplitude),length(AngleInp),length(SigmaKernel)]);
  
    for i = 1:length(Amplitude)
        i
        for j = 1:length(AngleInp)
            for s = 1:length(SigmaKernel)
                PPP = fspecial_XTHarmonicGaussian_new_copy('motion',Amplitude(i),AngleInp(j),SigmaKernel(s),upSample);
%                 deconvImage  = deconvlucy(proj1,scanKernel,bestIteration-1,0,WT);
%                 errorMatrix(i,j,s) = corr2(My_stripzero( deconvImage.*after_conv_small,CORR2Size),My_stripzero( simuShift.*after_conv_small,CORR2Size));  %IMPORTANT: has to be masked
                scanKernel = My_paddzero(PPP,[31 31]);
                CORR2test = zeros(1,15);
                            
                    % find best iteration
                    deconvImage=proj1.*after_conv_small;
                    CORR2test(1) = corr2(My_stripzero( deconvImage, CORR2Size),My_stripzero( simuShift.*after_conv_small,CORR2Size));

                    pj1deconv  = deconvlucy({proj1},scanKernel,1,0,WT);
                    deconvImage=pj1deconv{2}.*after_conv_small;
                    CORR2test(2) = corr2(My_stripzero( deconvImage, CORR2Size),My_stripzero( simuShift.*after_conv_small,CORR2Size));

                    for iter = 3:15
                        pj1deconv  = deconvlucy(pj1deconv,scanKernel,1,0,WT);
                        deconvImage=pj1deconv{2}.*after_conv_small;
                        CORR2test(iter) = corr2(My_stripzero( deconvImage,CORR2Size),My_stripzero( simuShift.*after_conv_small,CORR2Size));
                    end
                
                errorMatrix(i,j,s)  = max(CORR2test);
                        
            end
        end
    end
    
    [maxValue, scanInd]=max(errorMatrix(:));
    [AmpFinal, AngleFinal,SigmaFinal] = ind2sub(size(errorMatrix),scanInd);
    kScanBest(proj,:) = [Amplitude(AmpFinal),AngleInp(AngleFinal),SigmaKernel(SigmaFinal)];
    kScanBest(proj,:)

    PPPFinal = fspecial_XTHarmonicGaussian_new_copy('motion',kScanBest(proj,1),kScanBest(proj,2),kScanBest(proj,3),upSample);
    scanKernel = My_paddzero(PPPFinal,[31 31 ]);
    kernelVeryFinal = scanKernel;
    deconvImageFinal  = deconvlucy(proj1,kernelVeryFinal,bestIteration-1,0,WT);
    
    figure(12);
    subplot(121);imagesc(My_stripzero( deconvImageFinal.*after_conv_small,CORR2Size));axis image;title('deonvFinal');
    subplot(122);imagesc(My_stripzero(simuShift.*after_conv_small,CORR2Size));axis image;title('DIFF');pause(1)
    %% Final iteration test
    fprintf('FINAL ITERATION TEST\n');
    
    figure(33);subplot(122);imagesc(kernelVeryFinal);axis image;title('kernelLast')
    
    proj1 = (Projs(cropVectors(proj,1)+arr,cropVectors(proj,2)+arr,proj)-k(3))/k(7);
    WT = zeros(size(proj1));
    edge = 7;
    WT(edge:end-edge+1,edge:end-edge+1) = 1;
    
    deconvImage=proj1.*after_conv_small;
    CORR2ARFinal(proj,1) = corr2(My_stripzero( deconvImage, CORR2Size),My_stripzero( simuShift.*after_conv_small, CORR2Size));
    
    pj1deconv  = deconvlucy({proj1},kernelVeryFinal,1,0,WT);
    deconvImage=pj1deconv{2}.*after_conv_small;
    CORR2ARFinal(proj,2) = corr2(My_stripzero( deconvImage, CORR2Size),My_stripzero( simuShift.*after_conv_small,CORR2Size));
    
    for iter = 3:50
        pj1deconv  = deconvlucy(pj1deconv,kernelVeryFinal,1,0,WT);
        deconvImage=pj1deconv{2}.*after_conv_small;
        CORR2ARFinal(proj,iter) = corr2(My_stripzero( deconvImage, CORR2Size),My_stripzero( simuShift.*after_conv_small,CORR2Size));
    end
    
    figure(42);plot(CORR2ARFinal(proj,:),'r-o');title(num2str(proj));pause(2)
    
    %final deconvolution
    [maxValue, bestIteration]=max(CORR2ARFinal(proj,:));
    bestIteration
    
    if bestIteration==1
        pj1deconv = (Projs(:,:,proj)-k(3))/k(7);
    else
        WT = zeros(size(Projs(:,:,proj)));
        edge = 11;
        WT(edge:end-edge+1,edge:end-edge+1) = 1;
        pj1deconv  = deconvlucy((Projs(:,:,proj)-k(3))/k(7),kernelVeryFinal,bestIteration-1,0,WT);
    end
    DeconvFit(:,:,proj) = pj1deconv;
    
    loopShowTemp(pj1deconv);


end


%%

for proj = 1:size(DeconvFit,3)

    NewDset(:,:,proj) = My_binning_2b2_from3b3(DeconvFit(:,:,proj));
end

save('Projs_A_UP2Bin2.mat','NewDset')





